diff --git a/example/fuse_bn_act.py b/example/fuse_bn_act.py index 7d2bb76e..e9f740c1 100644 --- a/example/fuse_bn_act.py +++ b/example/fuse_bn_act.py @@ -1,11 +1,11 @@ # Copyright 2021 Yan Yan -# +# # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at -# +# # http://www.apache.org/licenses/LICENSE-2.0 -# +# # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. @@ -27,7 +27,7 @@ from spconv.core import ConvAlgo import torch.fx import spconv.pytorch as spconv -import copy +import copy import pickle from spconv.pytorch.conv import SparseConvolution from spconv.pytorch import functional as Fsp @@ -36,7 +36,11 @@ def fuse_bn_weights(conv_w_OKI, conv_b, bn_rm, bn_rv, bn_eps, bn_w, bn_b): NDim = conv_w_OKI.ndim - 2 permute = [0, NDim+1] + [i+1 for i in range(NDim)] - conv_w_OIK = conv_w_OKI.permute(*permute) + if list(conv_w_OKI.shape[1:-1]) != [1,1,1]: + conv_w_OIK = conv_w_OKI.permute(*permute) + else: + conv_w_OIK = conv_w_OKI.reshape((conv_w_OKI.shape[-1], conv_w_OKI.shape[0])).T.reshape(conv_w_OKI.shape).permute(*permute) + # OIDHW if conv_b is None: conv_b = torch.zeros_like(bn_rm) @@ -49,7 +53,11 @@ def fuse_bn_weights(conv_w_OKI, conv_b, bn_rm, bn_rv, bn_eps, bn_w, bn_b): conv_w_OIK = conv_w_OIK * (bn_w * bn_var_rsqrt).reshape([-1] + [1] * (len(conv_w_OIK.shape) - 1)) conv_b = (conv_b - bn_rm) * bn_var_rsqrt * bn_w + bn_b permute = [0,] + [i+2 for i in range(NDim)] + [1,] - conv_w_OKI = conv_w_OIK.permute(*permute).contiguous() + if list(conv_w_OKI.shape[1:-1]) != [1,1,1]: + conv_w_OKI = conv_w_OIK.permute(*permute).contiguous() + else: + conv_w_OKI = conv_w_OIK.permute(*permute).contiguous().squeeze().T.reshape(conv_w_OKI.shape) + return torch.nn.Parameter(conv_w_OKI), torch.nn.Parameter(conv_b) def fuse_bn(conv, bn): @@ -144,7 +152,7 @@ def fuse(model: torch.fx.GraphModule) -> torch.fx.GraphModule: # Now that all uses of the batch norm have been replaced, we can # safely remove the batch norm. fx_model.graph.erase_node(node) - + fx_model.graph.lint() # After we've modified our graph, we need to recompile our graph in order # to keep the generated code in sync. @@ -193,7 +201,7 @@ def fuse_act(model: torch.fx.GraphModule) -> torch.fx.GraphModule: # Now that all uses of the batch norm have been replaced, we can # safely remove the batch norm. fx_model.graph.erase_node(node) - + fx_model.graph.lint() # After we've modified our graph, we need to recompile our graph in order # to keep the generated code in sync. @@ -360,6 +368,12 @@ def __init__(self, shape, algo): nn.BatchNorm1d(64), nn.ReLU(), + spconv.SparseConv3d(64, + 64, + 1, + bias=False, + algo=algo), + nn.BatchNorm1d(64), ) max_batch_size = 1 # grid (dense map) is used for indice generation. use pre-allocated grid can run faster. @@ -397,7 +411,8 @@ def main(): device = torch.device("cuda:0") device_cpu = torch.device("cpu:0") dtype = torch.float32 - net = Net(spatial_shape, ConvAlgo.MaskImplicitGemm).cuda().eval().to(dtype) + net = Net(spatial_shape, ConvAlgo.Native).cuda().eval().to(dtype) + net.eval() tracer = MyTracer() graph_trace = tracer.trace(net) net_fused = torch.fx.GraphModule(tracer.root, graph_trace) @@ -414,10 +429,10 @@ def main(): out_fused = net_fused(voxels_th_cuda, coors_th_cuda, 1) res = Fsp.sparse_add_hash_based(out_ref, out_fused.minus()) print(torch.linalg.norm(res.features)) - _set_enable_int8_test_inplace(net_fused, True) - qvoxels_cuda = voxels_th_cuda.to(torch.int8) - - out_int8 = net_fused(qvoxels_cuda, coors_th_cuda, 1) + #_set_enable_int8_test_inplace(net_fused, True) + #qvoxels_cuda = voxels_th_cuda.to(torch.int8) + + #out_int8 = net_fused(qvoxels_cuda, coors_th_cuda, 1) if __name__ == "__main__": - main() \ No newline at end of file + main()