def __init__(self, in_nchannel, out_nchannel, D): ME.MinkowskiNetwork.__init__(self, D) channels = [in_nchannel, 16, 32] self.net = nn.Sequential( ME.MinkowskiStackSum( ME.MinkowskiConvolution( channels[0], channels[1], kernel_size=3, stride=1, dimension=D, ), nn.Sequential( ME.MinkowskiConvolution( channels[0], channels[1], kernel_size=3, stride=2, dimension=D, ), ME.MinkowskiStackSum( nn.Identity(), nn.Sequential( ME.MinkowskiConvolution( channels[1], channels[2], kernel_size=3, stride=2, dimension=D, ), ME.MinkowskiConvolutionTranspose( channels[2], channels[1], kernel_size=3, stride=1, dimension=D, ), ME.MinkowskiPoolingTranspose( kernel_size=2, stride=2, dimension=D ), ), ), ME.MinkowskiPoolingTranspose(kernel_size=2, stride=2, dimension=D), ), ), ME.MinkowskiToFeature(), nn.Linear(channels[1], out_nchannel, bias=True), )
def test_sum(self): coords, colors, pcd = load_file("1.ply") device = "cuda" D = 3 batch_size = 16 voxel_size = 0.02 channels = [3, 64, 128] dcoords = torch.from_numpy(np.floor(coords / voxel_size)).int() bcoords = batched_coordinates([dcoords for i in range(batch_size)]) in_feats = torch.rand(len(bcoords), 3).to(0) layer = MinkowskiStackSum( ME.MinkowskiConvolution( channels[0], channels[1], kernel_size=3, stride=1, dimension=3, ), nn.Sequential( ME.MinkowskiConvolution( channels[0], channels[1], kernel_size=3, stride=2, dimension=3, ), ME.MinkowskiStackSum( nn.Identity(), nn.Sequential( ME.MinkowskiConvolution( channels[1], channels[2], kernel_size=3, stride=2, dimension=3, ), ME.MinkowskiConvolutionTranspose( channels[2], channels[1], kernel_size=3, stride=1, dimension=3, ), ME.MinkowskiPoolingTranspose( kernel_size=2, stride=2, dimension=D ), ), ), ME.MinkowskiPoolingTranspose(kernel_size=2, stride=2, dimension=D), ), ).cuda() for i in range(1000): torch.cuda.empty_cache() sinput = ME.SparseTensor(in_feats, coordinates=bcoords, device=device) layer(sinput)