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)
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 conv_on_coords(): in_channels, out_channels, D = 2, 3, 2 coords, feats, labels = data_loader(in_channels, batch_size=1) # Create input with tensor stride == 4 strided_coords4, tensor_stride4 = get_random_coords(tensor_stride=4) strided_coords2, tensor_stride2 = get_random_coords(tensor_stride=2) input = ME.SparseTensor( feats=torch.rand(len(strided_coords4), in_channels), # coords=strided_coords4, tensor_stride=tensor_stride4) cm = input.coords_man # Convolution transpose and generate new coordinates conv_tr = ME.MinkowskiConvolutionTranspose(in_channels, out_channels, kernel_size=3, stride=2, has_bias=False, dimension=D) pool_tr = ME.MinkowskiPoolingTranspose(in_channels, out_channels, kernel_size=2, stride=2, has_bias=False, dimension=D) # If the there is no coordinates defined for the tensor stride, it will create one # tensor stride 4 -> conv_tr with stride 2 -> tensor stride 2 output1 = conv_tr(input) # output1 = pool_tr(input) # convolution on the specified coords output2 = conv_tr(input, coords) # output2 = pool_tr(input, coords) # convolution on the specified coords with tensor stride == 2 coords_key = cm.create_coords_key(strided_coords2, tensor_stride=2) output3 = conv_tr(input, coords_key) # output3 = pool_tr(input, coords_key) # convolution on the coordinates of a sparse tensor output4 = conv_tr(input, output1)
def __init__(self, in_channels, out_channels, config, D=3, **kwargs): super(FCNet, self).__init__(in_channels, out_channels, config, D) reps = 2 # Output of the first conv concated to conv6 self.conv1p1s1 = ME.MinkowskiConvolution(in_channels=in_channels, out_channels=self.PLANES[0], pixel_dist=1, kernel_size=3, has_bias=False, dimension=D) self.bn1 = nn.BatchNorm1d(self.PLANES[0]) self.block1 = FCNBlocks(self.PLANES[0], pixel_dist=1, reps=reps, D=D) self.conv2p1s2 = ME.MinkowskiConvolution(in_channels=self.PLANES[0], out_channels=self.PLANES[1], pixel_dist=1, kernel_size=2, stride=2, has_bias=False, dimension=D) self.bn2 = nn.BatchNorm1d(self.PLANES[1]) self.block2 = FCNBlocks(self.PLANES[1], pixel_dist=2, reps=reps, D=D) self.unpool2 = ME.MinkowskiPoolingTranspose(pixel_dist=2, kernel_size=2, stride=2, dimension=D) self.conv3p2s2 = ME.MinkowskiConvolution(in_channels=self.PLANES[1], out_channels=self.PLANES[2], pixel_dist=2, kernel_size=2, stride=2, has_bias=False, dimension=D) self.bn3 = nn.BatchNorm1d(self.PLANES[2]) self.block3 = FCNBlocks(self.PLANES[2], pixel_dist=4, reps=reps, D=D) self.unpool3 = ME.MinkowskiPoolingTranspose(pixel_dist=4, kernel_size=4, stride=4, dimension=D) self.conv4p4s2 = ME.MinkowskiConvolution(in_channels=self.PLANES[2], out_channels=self.PLANES[3], pixel_dist=4, kernel_size=2, stride=2, has_bias=False, dimension=D) self.bn4 = nn.BatchNorm1d(self.PLANES[3]) self.block4 = FCNBlocks(self.PLANES[3], pixel_dist=8, reps=reps, D=D) self.unpool4 = ME.MinkowskiPoolingTranspose(pixel_dist=8, kernel_size=8, stride=8, dimension=D) self.relu = nn.ReLU(inplace=True) self.final = ME.MinkowskiConvolution(in_channels=sum(self.PLANES[:4]), out_channels=out_channels, pixel_dist=1, kernel_size=1, stride=1, has_bias=True, dimension=D)
def __init__(self, in_channels=3, out_channels=32, bn_momentum=0.1, conv1_kernel_size=3, normalize_feature=False, D=3): ME.MinkowskiNetwork.__init__(self, D) NORM_TYPE = self.NORM_TYPE BLOCK_NORM_TYPE = self.BLOCK_NORM_TYPE CHANNELS = self.CHANNELS TR_CHANNELS = self.TR_CHANNELS DEPTHS = self.DEPTHS REGION_TYPE = self.REGION_TYPE self.normalize_feature = normalize_feature self.conv1 = conv( in_channels=in_channels, out_channels=CHANNELS[1], kernel_size=conv1_kernel_size, stride=1, dilation=1, has_bias=False, region_type=REGION_TYPE, dimension=D) self.norm1 = get_norm(NORM_TYPE, CHANNELS[1], bn_momentum=bn_momentum, dimension=D) self.block1 = nn.Sequential(*[ get_block( BLOCK_NORM_TYPE, CHANNELS[1], CHANNELS[1], bn_momentum=bn_momentum, region_type=REGION_TYPE, dimension=D) for d in range(DEPTHS[1]) ]) self.pool2 = ME.MinkowskiSumPooling(kernel_size=2, stride=2, dimension=D) self.conv2 = conv( in_channels=CHANNELS[1], out_channels=CHANNELS[2], kernel_size=1, stride=1, dilation=1, has_bias=False, dimension=D) self.norm2 = get_norm(NORM_TYPE, CHANNELS[2], bn_momentum=bn_momentum, dimension=D) self.block2 = nn.Sequential(*[ get_block( BLOCK_NORM_TYPE, CHANNELS[2], CHANNELS[2], bn_momentum=bn_momentum, region_type=REGION_TYPE, dimension=D) for d in range(DEPTHS[2]) ]) self.pool3 = ME.MinkowskiSumPooling(kernel_size=2, stride=2, dimension=D) self.conv3 = conv( in_channels=CHANNELS[2], out_channels=CHANNELS[3], kernel_size=1, stride=1, dilation=1, has_bias=False, dimension=D) self.norm3 = get_norm(NORM_TYPE, CHANNELS[3], bn_momentum=bn_momentum, dimension=D) self.block3 = nn.Sequential(*[ get_block( BLOCK_NORM_TYPE, CHANNELS[3], CHANNELS[3], bn_momentum=bn_momentum, region_type=REGION_TYPE, dimension=D) for d in range(DEPTHS[3]) ]) self.pool3_tr = ME.MinkowskiPoolingTranspose(kernel_size=2, stride=2, dimension=D) self.conv3_tr = conv_tr( in_channels=CHANNELS[3], out_channels=TR_CHANNELS[3], kernel_size=1, stride=1, dilation=1, has_bias=False, dimension=D) self.norm3_tr = get_norm( NORM_TYPE, TR_CHANNELS[3], bn_momentum=bn_momentum, dimension=D) self.block3_tr = nn.Sequential(*[ get_block( BLOCK_NORM_TYPE, TR_CHANNELS[3], TR_CHANNELS[3], bn_momentum=bn_momentum, region_type=REGION_TYPE, dimension=D) for d in range(DEPTHS[-3]) ]) self.pool2_tr = ME.MinkowskiPoolingTranspose(kernel_size=2, stride=2, dimension=D) self.conv2_tr = conv_tr( in_channels=CHANNELS[2] + TR_CHANNELS[3], out_channels=TR_CHANNELS[2], kernel_size=1, stride=1, dilation=1, has_bias=False, region_type=REGION_TYPE, dimension=D) self.norm2_tr = get_norm( NORM_TYPE, TR_CHANNELS[2], bn_momentum=bn_momentum, dimension=D) self.block2_tr = nn.Sequential(*[ get_block( BLOCK_NORM_TYPE, TR_CHANNELS[2], TR_CHANNELS[2], bn_momentum=bn_momentum, region_type=REGION_TYPE, dimension=D) for d in range(DEPTHS[-2]) ]) self.conv1_tr = conv_tr( in_channels=CHANNELS[1] + TR_CHANNELS[2], out_channels=TR_CHANNELS[1], kernel_size=1, stride=1, dilation=1, has_bias=False, region_type=REGION_TYPE, dimension=D) # self.block1_tr = BasicBlockBN(TR_CHANNELS[1], TR_CHANNELS[1], bn_momentum=bn_momentum, dimension=D) self.final = conv( in_channels=TR_CHANNELS[1], out_channels=out_channels, kernel_size=1, stride=1, dilation=1, has_bias=True, dimension=D)
def network_initialization(self, in_channels, out_channels, config, D): # Setup net_metadata dilations = self.DILATIONS bn_momentum = config.bn_momentum def space_n_time_m(n, m): return n if D == 3 else [n, n, n, m] if D == 4: self.OUT_PIXEL_DIST = space_n_time_m(self.OUT_PIXEL_DIST, 1) # Output of the first conv concated to conv6 self.inplanes = self.INIT_DIM self.conv1p1s1 = conv(in_channels, self.inplanes, kernel_size=space_n_time_m( config.conv1_kernel_size, 1), stride=1, dilation=1, conv_type=self.NON_BLOCK_CONV_TYPE, D=D) self.bn1 = get_norm(self.NORM_TYPE, self.PLANES[0], D, bn_momentum=bn_momentum) self.block1 = self._make_layer(self.BLOCK, self.PLANES[0], self.LAYERS[0], dilation=dilations[0], norm_type=self.NORM_TYPE, bn_momentum=bn_momentum) self.conv2p1s2 = conv(self.inplanes, self.inplanes, kernel_size=space_n_time_m(2, 1), stride=space_n_time_m(2, 1), dilation=1, conv_type=self.NON_BLOCK_CONV_TYPE, D=D) self.bn2 = get_norm(self.NORM_TYPE, self.inplanes, D, bn_momentum=bn_momentum) self.block2 = self._make_layer(self.BLOCK, self.PLANES[1], self.LAYERS[1], dilation=dilations[1], norm_type=self.NORM_TYPE, bn_momentum=bn_momentum) self.conv3p2s2 = conv(self.inplanes, self.inplanes, kernel_size=space_n_time_m(2, 1), stride=space_n_time_m(2, 1), dilation=1, conv_type=self.NON_BLOCK_CONV_TYPE, D=D) self.bn3 = get_norm(self.NORM_TYPE, self.inplanes, D, bn_momentum=bn_momentum) self.block3 = self._make_layer(self.BLOCK, self.PLANES[2], self.LAYERS[2], dilation=dilations[2], norm_type=self.NORM_TYPE, bn_momentum=bn_momentum) self.conv4p4s2 = conv(self.inplanes, self.inplanes, kernel_size=space_n_time_m(2, 1), stride=space_n_time_m(2, 1), dilation=1, conv_type=self.NON_BLOCK_CONV_TYPE, D=D) self.bn4 = get_norm(self.NORM_TYPE, self.inplanes, D, bn_momentum=bn_momentum) self.block4 = self._make_layer(self.BLOCK, self.PLANES[3], self.LAYERS[3], dilation=dilations[3], norm_type=self.NORM_TYPE, bn_momentum=bn_momentum) self.pool_tr4 = ME.MinkowskiPoolingTranspose(kernel_size=8, stride=8, dimension=D) out_pool4 = self.inplanes self.convtr4p8s2 = conv_tr(self.inplanes, self.PLANES[4], kernel_size=space_n_time_m(2, 1), upsample_stride=space_n_time_m(2, 1), dilation=1, bias=False, conv_type=self.NON_BLOCK_CONV_TYPE, D=D) self.bntr4 = get_norm(self.NORM_TYPE, self.PLANES[4], D, bn_momentum=bn_momentum) self.inplanes = self.PLANES[4] + self.PLANES[2] * self.BLOCK.expansion self.block5 = self._make_layer(self.BLOCK, self.PLANES[4], self.LAYERS[4], dilation=dilations[4], norm_type=self.NORM_TYPE, bn_momentum=bn_momentum) self.pool_tr5 = ME.MinkowskiPoolingTranspose(kernel_size=4, stride=4, dimension=D) out_pool5 = self.inplanes self.convtr5p4s2 = conv_tr(self.inplanes, self.PLANES[5], kernel_size=space_n_time_m(2, 1), upsample_stride=space_n_time_m(2, 1), dilation=1, bias=False, conv_type=self.NON_BLOCK_CONV_TYPE, D=D) self.bntr5 = get_norm(self.NORM_TYPE, self.PLANES[5], D, bn_momentum=bn_momentum) self.inplanes = self.PLANES[5] + self.PLANES[1] * self.BLOCK.expansion self.block6 = self._make_layer(self.BLOCK, self.PLANES[5], self.LAYERS[5], dilation=dilations[5], norm_type=self.NORM_TYPE, bn_momentum=bn_momentum) self.pool_tr6 = ME.MinkowskiPoolingTranspose(kernel_size=2, stride=2, dimension=D) out_pool6 = self.inplanes self.convtr6p2s2 = conv_tr(self.inplanes, self.PLANES[6], kernel_size=space_n_time_m(2, 1), upsample_stride=space_n_time_m(2, 1), dilation=1, bias=False, conv_type=self.NON_BLOCK_CONV_TYPE, D=D) self.bntr6 = get_norm(self.NORM_TYPE, self.PLANES[6], D, bn_momentum=bn_momentum) self.relu = MinkowskiReLU(inplace=True) self.final = nn.Sequential( conv(out_pool5 + out_pool6 + self.PLANES[6] + self.PLANES[0] * self.BLOCK.expansion, 512, kernel_size=1, bias=False, D=D), ME.MinkowskiBatchNorm(512), ME.MinkowskiReLU(), conv(512, out_channels, kernel_size=1, bias=True, D=D))