def __init__(self, in_nchannel, out_nchannel, D): super(UNet, self).__init__(D) self.block1 = torch.nn.Sequential( ME.MinkowskiConvolution( in_channels=in_nchannel, out_channels=8, kernel_size=3, stride=1, dimension=D), ME.MinkowskiBatchNorm(8)) self.block2 = torch.nn.Sequential( ME.MinkowskiConvolution( in_channels=8, out_channels=16, kernel_size=3, stride=2, dimension=D), ME.MinkowskiBatchNorm(16), ) self.block3 = torch.nn.Sequential( ME.MinkowskiConvolution( in_channels=16, out_channels=32, kernel_size=3, stride=2, dimension=D), ME.MinkowskiBatchNorm(32)) self.block3_tr = torch.nn.Sequential( ME.MinkowskiConvolutionTranspose( in_channels=32, out_channels=16, kernel_size=3, stride=2, dimension=D), ME.MinkowskiBatchNorm(16)) self.block2_tr = torch.nn.Sequential( ME.MinkowskiConvolutionTranspose( in_channels=32, out_channels=16, kernel_size=3, stride=2, dimension=D), ME.MinkowskiBatchNorm(16)) self.conv1_tr = ME.MinkowskiConvolution( in_channels=24, out_channels=out_nchannel, kernel_size=1, stride=1, dimension=D)
def network_initialization(self, in_channels, config, D): up_kernel_size = 3 self.conv_up1 = nn.Sequential( ME.MinkowskiConvolutionTranspose( in_channels[0], in_channels[0], kernel_size=up_kernel_size, stride=2, generate_new_coords=True, dimension=3), ME.MinkowskiBatchNorm(in_channels[0]), ME.MinkowskiELU()) self.conv_up2 = nn.Sequential( ME.MinkowskiConvolutionTranspose( in_channels[1], in_channels[0], kernel_size=up_kernel_size, stride=2, generate_new_coords=True, dimension=3), ME.MinkowskiBatchNorm(in_channels[0]), ME.MinkowskiELU()) self.conv_up3 = nn.Sequential( ME.MinkowskiConvolutionTranspose( in_channels[2], in_channels[1], kernel_size=up_kernel_size, stride=2, generate_new_coords=True, dimension=3), ME.MinkowskiBatchNorm(in_channels[1]), ME.MinkowskiELU()) self.conv_up4 = nn.Sequential( ME.MinkowskiConvolutionTranspose( in_channels[3], in_channels[2], kernel_size=up_kernel_size, stride=2, generate_new_coords=True, dimension=3), ME.MinkowskiBatchNorm(in_channels[2]), ME.MinkowskiELU()) self.conv_feat1 = nn.Sequential( ME.MinkowskiConvolution( in_channels[0], config.upsample_feat_size, kernel_size=1, dimension=3), ME.MinkowskiBatchNorm(config.upsample_feat_size), ME.MinkowskiELU()) self.conv_feat2 = nn.Sequential( ME.MinkowskiConvolution( in_channels[1], config.upsample_feat_size, kernel_size=1, dimension=3), ME.MinkowskiBatchNorm(config.upsample_feat_size), ME.MinkowskiELU()) self.conv_feat3 = nn.Sequential( ME.MinkowskiConvolution( in_channels[2], config.upsample_feat_size, kernel_size=1, dimension=3), ME.MinkowskiBatchNorm(config.upsample_feat_size), ME.MinkowskiELU()) self.conv_feat4 = nn.Sequential( ME.MinkowskiConvolution( in_channels[3], config.upsample_feat_size, kernel_size=1, dimension=3), ME.MinkowskiBatchNorm(config.upsample_feat_size), ME.MinkowskiELU())
def __init__(self, input_a_dim, input_b_dim, out_dim, kernel_size=2): super().__init__() ''' Deconv x_a concat with x_b then apply output-projection ''' self.input_a_dim = input_a_dim self.input_b_dim = input_b_dim self.out_dim = out_dim self.conv_a = nn.Sequential( ME.MinkowskiConvolutionTranspose(in_channels=input_a_dim, out_channels=input_a_dim, kernel_size=4, stride=4, dimension=3), ME.MinkowskiBatchNorm(input_a_dim), ME.MinkowskiReLU(), ) self.conv_proj = nn.Sequential( ME.MinkowskiConvolution(in_channels=input_a_dim + input_b_dim, out_channels=out_dim, kernel_size=3, stride=1, dimension=3), ME.MinkowskiBatchNorm(out_dim), ME.MinkowskiReLU(), )
def conv_tr(in_planes, out_planes, kernel_size, upsample_stride=1, dilation=1, bias=False, conv_type=ConvType.HYPERCUBE, D=-1): assert D > 0, 'Dimension must be a positive integer' region_type, axis_types, kernel_size = convert_conv_type(conv_type, kernel_size, D) kernel_generator = ME.KernelGenerator( kernel_size, upsample_stride, dilation, region_type=region_type, axis_types=axis_types, dimension=D) return ME.MinkowskiConvolutionTranspose( in_channels=in_planes, out_channels=out_planes, kernel_size=kernel_size, stride=upsample_stride, dilation=dilation, has_bias=bias, kernel_generator=kernel_generator, dimension=D)
def __init__(self, up_conv_nn=[], kernel_size=3, stride=1, dilation=1, has_bias=False, activation=ME.MinkowskiReLU(inplace=True), bn_momentum=0.01, dimension=-1, **kwargs): """ Block convolution which consists of a convolution a batch norm a block operation and an activation. the block operation is usually a resnetBlock """ # instantiate convolution # instantiate batchnorm # instantiate block # activation super(SimpleBlockUp, self).__init__() self.conv_tr = ME.MinkowskiConvolutionTranspose( up_conv_nn[0], up_conv_nn[1], kernel_size=kernel_size, stride=stride, dilation=dilation, dimension=dimension) self.bn = ME.MinkowskiBatchNorm(up_conv_nn[1], momentum=bn_momentum) self.block = BasicBlock(up_conv_nn[1], up_conv_nn[1], bn_momentum=bn_momentum, dimension=dimension) self.activation = activation
def conv_tr(in_channels, out_channels, kernel_size, stride=1, dilation=1, bias=False, region_type=ME.RegionType.HYPER_CUBE, dimension=-1): assert dimension > 0, 'Dimension must be a positive integer' kernel_generator = ME.KernelGenerator(kernel_size, stride, dilation, is_transpose=True, region_type=region_type, dimension=dimension) kernel_generator = ME.KernelGenerator(kernel_size, stride, dilation, is_transpose=True, region_type=region_type, dimension=dimension) return ME.MinkowskiConvolutionTranspose(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, dilation=dilation, bias=bias, kernel_generator=kernel_generator, dimension=dimension)
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, inc, outc, ks=3, stride=1, D=3): super(BasicDeconvolutionBlock, self).__init__() self.net = nn.Sequential( ME.MinkowskiConvolutionTranspose( inc, outc, kernel_size=ks, stride=stride, dimension=D), ME.MinkowskiBatchNorm(outc), ME.MinkowskiReLU(True)) nn.init.constant_(self.net[1].bn.weight, 1.0) nn.init.constant_(self.net[1].bn.bias, 0.0)
def __init__(self, in_nchannel, out_nchannel, D): super(UNet, self).__init__(D) self.conv1 = ME.MinkowskiConvolution( in_channels=in_nchannel, out_channels=8, kernel_size=3, stride=1, dimension=D) self.bn1 = ME.MinkowskiBatchNorm(8) self.conv2 = ME.MinkowskiConvolution( in_channels=8, out_channels=16, kernel_size=3, stride=2, dimension=D) self.bn2 = ME.MinkowskiBatchNorm(16) self.conv3 = ME.MinkowskiConvolution( in_channels=16, out_channels=32, kernel_size=3, stride=2, dimension=D) self.bn3 = ME.MinkowskiBatchNorm(32) self.conv4 = ME.MinkowskiConvolutionTranspose( in_channels=32, out_channels=16, kernel_size=3, stride=2, dimension=D) self.bn4 = ME.MinkowskiBatchNorm(16) self.conv5 = ME.MinkowskiConvolutionTranspose( in_channels=32, out_channels=16, kernel_size=3, stride=2, dimension=D) self.bn5 = ME.MinkowskiBatchNorm(16) self.conv6 = ME.MinkowskiConvolution( in_channels=24, out_channels=out_nchannel, kernel_size=1, stride=1, dimension=D)
def __init__(self, inc, outc, ks=3, stride=1, D=3): super(BasicDeconvolutionBlock4d, self).__init__() self.net = nn.Sequential( ME.MinkowskiConvolutionTranspose( inc, outc, stride=stride, dimension=D, kernel_generator=ME.KernelGenerator( kernel_size=(ks, ks, ks, 1), dimension=D, region_type=ME.RegionType.HYPERCUBE)), ME.MinkowskiBatchNorm(outc), ME.MinkowskiReLU(True)) nn.init.constant_(self.net[1].bn.weight, 1.0) nn.init.constant_(self.net[1].bn.bias, 0.0)
def __init__(self, in_channels, out_channels, D=3): nn.Module.__init__(self) self.net = nn.Sequential( ME.MinkowskiConvolution(in_channels, 32, 3, dimension=D), ME.MinkowskiBatchNorm(32), ME.MinkowskiReLU(), ME.MinkowskiConvolution(32, 64, 3, stride=2, dimension=D), ME.MinkowskiBatchNorm(64), ME.MinkowskiReLU(), ME.MinkowskiConvolutionTranspose(64, 32, 3, stride=2, dimension=D), ME.MinkowskiBatchNorm(32), ME.MinkowskiReLU(), ME.MinkowskiConvolution(32, out_channels, kernel_size=1, dimension=D), )
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 __init__(self, down_channels, skip_channels, out_channels, D=3): super().__init__() self.deconv = nn.Sequential( ME.MinkowskiConvolutionTranspose(down_channels, down_channels, kernel_size=3, stride=2, dimension=D), ME.MinkowskiReLU(inplace=True), ) self.conv = nn.Sequential( ME.MinkowskiConvolution(skip_channels + down_channels, out_channels, kernel_size=3, dimension=D), ME.MinkowskiReLU(inplace=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 conv(): in_channels, out_channels, D = 2, 3, 2 bp() coords, feats, labels = data_loader(in_channels, batch_size=1) # Convolution input = ME.SparseTensor(feats=feats, coords=coords) conv = ME.MinkowskiConvolution(in_channels, out_channels, kernel_size=3, stride=2, has_bias=False, dimension=D) output = conv(input) print('Input:') print_sparse_tensor(input) print('Output:') print_sparse_tensor(output) # Convolution transpose and generate new coordinates strided_coords, tensor_stride = get_random_coords() bp() input = ME.SparseTensor( feats=torch.rand(len(strided_coords), in_channels), # coords=strided_coords, tensor_stride=tensor_stride) conv_tr = ME.MinkowskiConvolutionTranspose(in_channels, out_channels, kernel_size=3, stride=2, has_bias=False, dimension=D) output = conv_tr(input) print('\nInput:') print_sparse_tensor(input) print('Convolution Transpose Output:') print_sparse_tensor(output)
def __init__(self, up_conv_nn, kernel_size, stride, dilation, dimension=3, bn_momentum=0.01, norm_type=NormType.BATCH_NORM, block_norm_type=NormType.BATCH_NORM, **kwargs): ME.MinkowskiNetwork.__init__(self, dimension) self.conv = ME.MinkowskiConvolutionTranspose( in_channels=up_conv_nn[0], out_channels=up_conv_nn[1], kernel_size=kernel_size, stride=stride, dilation=dilation, has_bias=False, dimension=dimension, ) if len(up_conv_nn) == 3: self.final = ME.MinkowskiConvolution( in_channels=up_conv_nn[1], out_channels=up_conv_nn[2], kernel_size=kernel_size, stride=stride, dilation=dilation, has_bias=True, dimension=dimension, ) else: self.norm = get_norm(norm_type, up_conv_nn[1], bn_momentum=bn_momentum, D=dimension) self.block = get_block(block_norm_type, up_conv_nn[1], up_conv_nn[1], bn_momentum=bn_momentum, D=dimension) self.final = None
def __init__(self, input_dim, out_dim, k=3): super().__init__() ''' Transition Up Layer npoint: number of input points nsample: k in kNN, default 3 in_dim: feature dimension of the input feature x (output of the PCTLayer) out_dim: feature dimension of the TDLayer ''' self.k = k self.input_dim = input_dim self.out_dim = out_dim self.linear_1 = nn.Conv1d(input_dim, out_dim, 1) self.linear_2 = nn.Conv1d(out_dim, out_dim, 1) self.conv = ME.MinkowskiConvolutionTranspose(in_channels=input_dim, out_channels=out_dim, kernel_size=3, stride=2, dimension=3)
def __init__(self, up_conv_nn=[], kernel_sizes=[], strides=[], dilations=[], kernel_size=3, stride=1, dilation=1, norm_layer=ME.MinkowskiBatchNorm, activation=ME.MinkowskiReLU, bn_momentum=0.1, dimension=-1, up_stride=2, skip=True, **kwargs): self.skip = skip super(ResnetBlockUp, self).__init__( up_conv_nn[0], up_conv_nn[1], up_conv_nn[2], kernel_sizes=kernel_sizes, strides=strides, dilations=dilations, kernel_size=kernel_size, stride=stride, dilation=dilation, norm_layer=norm_layer, activation=activation, bn_momentum=bn_momentum, dimension=dimension, ) self.upsample = ME.MinkowskiConvolutionTranspose(up_conv_nn[0], up_conv_nn[2], kernel_size=2, stride=up_stride, dimension=dimension)
def network_initialization(self, in_channels, out_channels, D): # Output of the first conv concated to conv6 self.inplanes = self.INIT_DIM self.conv0p1s1 = ME.MinkowskiConvolution(in_channels, self.inplanes, kernel_size=5, dimension=D) self.bn0 = ME.MinkowskiBatchNorm(self.inplanes) self.conv1p1s2 = ME.MinkowskiConvolution(self.inplanes, self.inplanes, kernel_size=2, stride=2, dimension=D) self.bn1 = ME.MinkowskiBatchNorm(self.inplanes) self.block1 = self._make_layer(self.BLOCK, self.PLANES[0], self.LAYERS[0]) self.conv2p2s2 = ME.MinkowskiConvolution(self.inplanes, self.inplanes, kernel_size=2, stride=2, dimension=D) self.bn2 = ME.MinkowskiBatchNorm(self.inplanes) self.block2 = self._make_layer(self.BLOCK, self.PLANES[1], self.LAYERS[1]) self.conv3p4s2 = ME.MinkowskiConvolution(self.inplanes, self.inplanes, kernel_size=2, stride=2, dimension=D) self.bn3 = ME.MinkowskiBatchNorm(self.inplanes) self.block3 = self._make_layer(self.BLOCK, self.PLANES[2], self.LAYERS[2]) self.conv4p8s2 = ME.MinkowskiConvolution(self.inplanes, self.inplanes, kernel_size=2, stride=2, dimension=D) self.bn4 = ME.MinkowskiBatchNorm(self.inplanes) self.block4 = self._make_layer(self.BLOCK, self.PLANES[3], self.LAYERS[3]) self.convtr4p16s2 = ME.MinkowskiConvolutionTranspose(self.inplanes, self.PLANES[4], kernel_size=2, stride=2, dimension=D) self.bntr4 = ME.MinkowskiBatchNorm(self.PLANES[4]) self.inplanes = self.PLANES[4] + self.PLANES[2] * self.BLOCK.expansion self.block5 = self._make_layer(self.BLOCK, self.PLANES[4], self.LAYERS[4]) self.convtr5p8s2 = ME.MinkowskiConvolutionTranspose(self.inplanes, self.PLANES[5], kernel_size=2, stride=2, dimension=D) self.bntr5 = ME.MinkowskiBatchNorm(self.PLANES[5]) self.inplanes = self.PLANES[5] + self.PLANES[1] * self.BLOCK.expansion self.block6 = self._make_layer(self.BLOCK, self.PLANES[5], self.LAYERS[5]) self.convtr6p4s2 = ME.MinkowskiConvolutionTranspose(self.inplanes, self.PLANES[6], kernel_size=2, stride=2, dimension=D) self.bntr6 = ME.MinkowskiBatchNorm(self.PLANES[6]) self.inplanes = self.PLANES[6] + self.PLANES[0] * self.BLOCK.expansion self.block7 = self._make_layer(self.BLOCK, self.PLANES[6], self.LAYERS[6]) self.convtr7p2s2 = ME.MinkowskiConvolutionTranspose(self.inplanes, self.PLANES[7], kernel_size=2, stride=2, dimension=D) self.bntr7 = ME.MinkowskiBatchNorm(self.PLANES[7]) self.inplanes = self.PLANES[7] + self.INIT_DIM self.block8 = self._make_layer(self.BLOCK, self.PLANES[7], self.LAYERS[7]) if ME.__version__.split('.')[1] < '5': self.final = ME.MinkowskiConvolution(self.PLANES[7] * self.BLOCK.expansion, out_channels, kernel_size=1, has_bias=True, dimension=D) else: self.final = ME.MinkowskiConvolution(self.PLANES[7] * self.BLOCK.expansion, out_channels, kernel_size=1, bias=True, dimension=D) self.relu = ME.MinkowskiReLU(inplace=True)
def __init__(self, in_channels=3, out_channels=32, bn_momentum=0.1, normalize_feature=None, conv1_kernel_size=None, 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 self.normalize_feature = normalize_feature self.conv1 = ME.MinkowskiConvolution( in_channels=in_channels, out_channels=CHANNELS[1], kernel_size=conv1_kernel_size, stride=1, dilation=1, has_bias=False, dimension=D) self.norm1 = get_norm(NORM_TYPE, CHANNELS[1], bn_momentum=bn_momentum, D=D) self.block1 = get_block( BLOCK_NORM_TYPE, CHANNELS[1], CHANNELS[1], bn_momentum=bn_momentum, D=D) self.conv2 = ME.MinkowskiConvolution( in_channels=CHANNELS[1], out_channels=CHANNELS[2], kernel_size=3, stride=2, dilation=1, has_bias=False, dimension=D) self.norm2 = get_norm(NORM_TYPE, CHANNELS[2], bn_momentum=bn_momentum, D=D) self.block2 = get_block( BLOCK_NORM_TYPE, CHANNELS[2], CHANNELS[2], bn_momentum=bn_momentum, D=D) self.conv3 = ME.MinkowskiConvolution( in_channels=CHANNELS[2], out_channels=CHANNELS[3], kernel_size=3, stride=2, dilation=1, has_bias=False, dimension=D) self.norm3 = get_norm(NORM_TYPE, CHANNELS[3], bn_momentum=bn_momentum, D=D) self.block3 = get_block( BLOCK_NORM_TYPE, CHANNELS[3], CHANNELS[3], bn_momentum=bn_momentum, D=D) self.conv4 = ME.MinkowskiConvolution( in_channels=CHANNELS[3], out_channels=CHANNELS[4], kernel_size=3, stride=2, dilation=1, has_bias=False, dimension=D) self.norm4 = get_norm(NORM_TYPE, CHANNELS[4], bn_momentum=bn_momentum, D=D) self.block4 = get_block( BLOCK_NORM_TYPE, CHANNELS[4], CHANNELS[4], bn_momentum=bn_momentum, D=D) self.conv4_tr = ME.MinkowskiConvolutionTranspose( in_channels=CHANNELS[4], out_channels=TR_CHANNELS[4], kernel_size=3, stride=2, dilation=1, has_bias=False, dimension=D) self.norm4_tr = get_norm(NORM_TYPE, TR_CHANNELS[4], bn_momentum=bn_momentum, D=D) self.block4_tr = get_block( BLOCK_NORM_TYPE, TR_CHANNELS[4], TR_CHANNELS[4], bn_momentum=bn_momentum, D=D) self.conv3_tr = ME.MinkowskiConvolutionTranspose( in_channels=CHANNELS[3] + TR_CHANNELS[4], out_channels=TR_CHANNELS[3], kernel_size=3, stride=2, dilation=1, has_bias=False, dimension=D) self.norm3_tr = get_norm(NORM_TYPE, TR_CHANNELS[3], bn_momentum=bn_momentum, D=D) self.block3_tr = get_block( BLOCK_NORM_TYPE, TR_CHANNELS[3], TR_CHANNELS[3], bn_momentum=bn_momentum, D=D) self.conv2_tr = ME.MinkowskiConvolutionTranspose( in_channels=CHANNELS[2] + TR_CHANNELS[3], out_channels=TR_CHANNELS[2], kernel_size=3, stride=2, dilation=1, has_bias=False, dimension=D) self.norm2_tr = get_norm(NORM_TYPE, TR_CHANNELS[2], bn_momentum=bn_momentum, D=D) self.block2_tr = get_block( BLOCK_NORM_TYPE, TR_CHANNELS[2], TR_CHANNELS[2], bn_momentum=bn_momentum, D=D) self.conv1_tr = ME.MinkowskiConvolution( in_channels=CHANNELS[1] + TR_CHANNELS[2], out_channels=TR_CHANNELS[1], kernel_size=1, stride=1, dilation=1, has_bias=False, dimension=D) # self.block1_tr = BasicBlockBN(TR_CHANNELS[1], TR_CHANNELS[1], bn_momentum=bn_momentum, D=D) self.final = ME.MinkowskiConvolution( in_channels=TR_CHANNELS[1], out_channels=out_channels, kernel_size=1, stride=1, dilation=1, has_bias=True, dimension=D)
def __init__(self): nn.Module.__init__(self) # Input sparse tensor must have tensor stride 128. ch = self.CHANNELS # Block 1 self.block1 = nn.Sequential( ME.MinkowskiConvolutionTranspose(ch[0], ch[0], kernel_size=2, stride=2, generate_new_coords=True, dimension=3), ME.MinkowskiBatchNorm(ch[0]), ME.MinkowskiELU(), ME.MinkowskiConvolution(ch[0], ch[0], kernel_size=3, dimension=3), ME.MinkowskiBatchNorm(ch[0]), ME.MinkowskiELU(), ME.MinkowskiConvolutionTranspose(ch[0], ch[1], kernel_size=2, stride=2, generate_new_coords=True, dimension=3), ME.MinkowskiBatchNorm(ch[1]), ME.MinkowskiELU(), ME.MinkowskiConvolution(ch[1], ch[1], kernel_size=3, dimension=3), ME.MinkowskiBatchNorm(ch[1]), ME.MinkowskiELU(), ) self.block1_cls = ME.MinkowskiConvolution(ch[1], 1, kernel_size=1, has_bias=True, dimension=3) # Block 2 self.block2 = nn.Sequential( ME.MinkowskiConvolutionTranspose(ch[1], ch[2], kernel_size=2, stride=2, generate_new_coords=True, dimension=3), ME.MinkowskiBatchNorm(ch[2]), ME.MinkowskiELU(), ME.MinkowskiConvolution(ch[2], ch[2], kernel_size=3, dimension=3), ME.MinkowskiBatchNorm(ch[2]), ME.MinkowskiELU(), ) self.block2_cls = ME.MinkowskiConvolution(ch[2], 1, kernel_size=1, has_bias=True, dimension=3) # Block 3 self.block3 = nn.Sequential( ME.MinkowskiConvolutionTranspose(ch[2], ch[3], kernel_size=2, stride=2, generate_new_coords=True, dimension=3), ME.MinkowskiBatchNorm(ch[3]), ME.MinkowskiELU(), ME.MinkowskiConvolution(ch[3], ch[3], kernel_size=3, dimension=3), ME.MinkowskiBatchNorm(ch[3]), ME.MinkowskiELU(), ) self.block3_cls = ME.MinkowskiConvolution(ch[3], 1, kernel_size=1, has_bias=True, dimension=3) # Block 4 self.block4 = nn.Sequential( ME.MinkowskiConvolutionTranspose(ch[3], ch[4], kernel_size=2, stride=2, generate_new_coords=True, dimension=3), ME.MinkowskiBatchNorm(ch[4]), ME.MinkowskiELU(), ME.MinkowskiConvolution(ch[4], ch[4], kernel_size=3, dimension=3), ME.MinkowskiBatchNorm(ch[4]), ME.MinkowskiELU(), ) self.block4_cls = ME.MinkowskiConvolution(ch[4], 1, kernel_size=1, has_bias=True, dimension=3) # Block 5 self.block5 = nn.Sequential( ME.MinkowskiConvolutionTranspose(ch[4], ch[5], kernel_size=2, stride=2, generate_new_coords=True, dimension=3), ME.MinkowskiBatchNorm(ch[5]), ME.MinkowskiELU(), ME.MinkowskiConvolution(ch[5], ch[5], kernel_size=3, dimension=3), ME.MinkowskiBatchNorm(ch[5]), ME.MinkowskiELU(), ) self.block5_cls = ME.MinkowskiConvolution(ch[5], 1, kernel_size=1, has_bias=True, dimension=3) # Block 6 self.block6 = nn.Sequential( ME.MinkowskiConvolutionTranspose(ch[5], ch[6], kernel_size=2, stride=2, generate_new_coords=True, dimension=3), ME.MinkowskiBatchNorm(ch[6]), ME.MinkowskiELU(), ME.MinkowskiConvolution(ch[6], ch[6], kernel_size=3, dimension=3), ME.MinkowskiBatchNorm(ch[6]), ME.MinkowskiELU(), ) self.block6_cls = ME.MinkowskiConvolution(ch[6], 1, kernel_size=1, has_bias=True, dimension=3) # pruning self.pruning = ME.MinkowskiPruning()
def __init__(self, resolution, in_nchannel=512): nn.Module.__init__(self) self.resolution = resolution # Input sparse tensor must have tensor stride 128. enc_ch = self.ENC_CHANNELS dec_ch = self.DEC_CHANNELS # Encoder self.enc_block_s1 = nn.Sequential( ME.MinkowskiConvolution(1, enc_ch[0], kernel_size=3, stride=1, dimension=3), ME.MinkowskiBatchNorm(enc_ch[0]), ME.MinkowskiELU(), ) self.enc_block_s1s2 = nn.Sequential( ME.MinkowskiConvolution(enc_ch[0], enc_ch[1], kernel_size=2, stride=2, dimension=3), ME.MinkowskiBatchNorm(enc_ch[1]), ME.MinkowskiELU(), ME.MinkowskiConvolution(enc_ch[1], enc_ch[1], kernel_size=3, dimension=3), ME.MinkowskiBatchNorm(enc_ch[1]), ME.MinkowskiELU(), ) self.enc_block_s2s4 = nn.Sequential( ME.MinkowskiConvolution(enc_ch[1], enc_ch[2], kernel_size=2, stride=2, dimension=3), ME.MinkowskiBatchNorm(enc_ch[2]), ME.MinkowskiELU(), ME.MinkowskiConvolution(enc_ch[2], enc_ch[2], kernel_size=3, dimension=3), ME.MinkowskiBatchNorm(enc_ch[2]), ME.MinkowskiELU(), ) self.enc_block_s4s8 = nn.Sequential( ME.MinkowskiConvolution(enc_ch[2], enc_ch[3], kernel_size=2, stride=2, dimension=3), ME.MinkowskiBatchNorm(enc_ch[3]), ME.MinkowskiELU(), ME.MinkowskiConvolution(enc_ch[3], enc_ch[3], kernel_size=3, dimension=3), ME.MinkowskiBatchNorm(enc_ch[3]), ME.MinkowskiELU(), ) self.enc_block_s8s16 = nn.Sequential( ME.MinkowskiConvolution(enc_ch[3], enc_ch[4], kernel_size=2, stride=2, dimension=3), ME.MinkowskiBatchNorm(enc_ch[4]), ME.MinkowskiELU(), ME.MinkowskiConvolution(enc_ch[4], enc_ch[4], kernel_size=3, dimension=3), ME.MinkowskiBatchNorm(enc_ch[4]), ME.MinkowskiELU(), ) self.enc_block_s16s32 = nn.Sequential( ME.MinkowskiConvolution(enc_ch[4], enc_ch[5], kernel_size=2, stride=2, dimension=3), ME.MinkowskiBatchNorm(enc_ch[5]), ME.MinkowskiELU(), ME.MinkowskiConvolution(enc_ch[5], enc_ch[5], kernel_size=3, dimension=3), ME.MinkowskiBatchNorm(enc_ch[5]), ME.MinkowskiELU(), ) self.enc_block_s32s64 = nn.Sequential( ME.MinkowskiConvolution(enc_ch[5], enc_ch[6], kernel_size=2, stride=2, dimension=3), ME.MinkowskiBatchNorm(enc_ch[6]), ME.MinkowskiELU(), ME.MinkowskiConvolution(enc_ch[6], enc_ch[6], kernel_size=3, dimension=3), ME.MinkowskiBatchNorm(enc_ch[6]), ME.MinkowskiELU(), ) # Decoder self.dec_block_s64s32 = nn.Sequential( ME.MinkowskiConvolutionTranspose(enc_ch[6], dec_ch[5], kernel_size=4, stride=2, generate_new_coords=True, dimension=3), ME.MinkowskiBatchNorm(dec_ch[5]), ME.MinkowskiELU(), ME.MinkowskiConvolution(dec_ch[5], dec_ch[5], kernel_size=3, dimension=3), ME.MinkowskiBatchNorm(dec_ch[5]), ME.MinkowskiELU(), ) self.dec_s32_cls = ME.MinkowskiConvolution(dec_ch[5], 1, kernel_size=1, has_bias=True, dimension=3) self.dec_block_s32s16 = nn.Sequential( ME.MinkowskiConvolutionTranspose(enc_ch[5], dec_ch[4], kernel_size=2, stride=2, generate_new_coords=True, dimension=3), ME.MinkowskiBatchNorm(dec_ch[4]), ME.MinkowskiELU(), ME.MinkowskiConvolution(dec_ch[4], dec_ch[4], kernel_size=3, dimension=3), ME.MinkowskiBatchNorm(dec_ch[4]), ME.MinkowskiELU(), ) self.dec_s16_cls = ME.MinkowskiConvolution(dec_ch[4], 1, kernel_size=1, has_bias=True, dimension=3) self.dec_block_s16s8 = nn.Sequential( ME.MinkowskiConvolutionTranspose(dec_ch[4], dec_ch[3], kernel_size=2, stride=2, generate_new_coords=True, dimension=3), ME.MinkowskiBatchNorm(dec_ch[3]), ME.MinkowskiELU(), ME.MinkowskiConvolution(dec_ch[3], dec_ch[3], kernel_size=3, dimension=3), ME.MinkowskiBatchNorm(dec_ch[3]), ME.MinkowskiELU(), ) self.dec_s8_cls = ME.MinkowskiConvolution(dec_ch[3], 1, kernel_size=1, has_bias=True, dimension=3) self.dec_block_s8s4 = nn.Sequential( ME.MinkowskiConvolutionTranspose(dec_ch[3], dec_ch[2], kernel_size=2, stride=2, generate_new_coords=True, dimension=3), ME.MinkowskiBatchNorm(dec_ch[2]), ME.MinkowskiELU(), ME.MinkowskiConvolution(dec_ch[2], dec_ch[2], kernel_size=3, dimension=3), ME.MinkowskiBatchNorm(dec_ch[2]), ME.MinkowskiELU(), ) self.dec_s4_cls = ME.MinkowskiConvolution(dec_ch[2], 1, kernel_size=1, has_bias=True, dimension=3) self.dec_block_s4s2 = nn.Sequential( ME.MinkowskiConvolutionTranspose(dec_ch[2], dec_ch[1], kernel_size=2, stride=2, generate_new_coords=True, dimension=3), ME.MinkowskiBatchNorm(dec_ch[1]), ME.MinkowskiELU(), ME.MinkowskiConvolution(dec_ch[1], dec_ch[1], kernel_size=3, dimension=3), ME.MinkowskiBatchNorm(dec_ch[1]), ME.MinkowskiELU(), ) self.dec_s2_cls = ME.MinkowskiConvolution(dec_ch[1], 1, kernel_size=1, has_bias=True, dimension=3) self.dec_block_s2s1 = nn.Sequential( ME.MinkowskiConvolutionTranspose(dec_ch[1], dec_ch[0], kernel_size=2, stride=2, generate_new_coords=True, dimension=3), ME.MinkowskiBatchNorm(dec_ch[0]), ME.MinkowskiELU(), ME.MinkowskiConvolution(dec_ch[0], dec_ch[0], kernel_size=3, dimension=3), ME.MinkowskiBatchNorm(dec_ch[0]), ME.MinkowskiELU(), ) self.dec_s1_cls = ME.MinkowskiConvolution(dec_ch[0], 1, kernel_size=1, has_bias=True, dimension=3) # pruning self.pruning = ME.MinkowskiPruning()
def __init__(self, in_channels=3, out_channels=32, bn_momentum=0.1, normalize_feature=None, conv1_kernel_size=None, D=3): super(SimpleNet, self).__init__(D) NORM_TYPE = self.NORM_TYPE CHANNELS = self.CHANNELS TR_CHANNELS = self.TR_CHANNELS self.normalize_feature = normalize_feature self.conv1 = ME.MinkowskiConvolution( in_channels=in_channels, out_channels=CHANNELS[1], kernel_size=conv1_kernel_size, stride=1, dilation=1, has_bias=False, dimension=D) self.norm1 = get_norm(NORM_TYPE, CHANNELS[1], bn_momentum=bn_momentum, D=D) self.conv2 = ME.MinkowskiConvolution( in_channels=CHANNELS[1], out_channels=CHANNELS[2], kernel_size=3, stride=2, dilation=1, has_bias=False, dimension=D) self.norm2 = get_norm(NORM_TYPE, CHANNELS[2], bn_momentum=bn_momentum, D=D) self.conv3 = ME.MinkowskiConvolution( in_channels=CHANNELS[2], out_channels=CHANNELS[3], kernel_size=3, stride=2, dilation=1, has_bias=False, dimension=D) self.norm3 = get_norm(NORM_TYPE, CHANNELS[3], bn_momentum=bn_momentum, D=D) self.conv3_tr = ME.MinkowskiConvolutionTranspose( in_channels=CHANNELS[3], out_channels=TR_CHANNELS[3], kernel_size=3, stride=2, dilation=1, has_bias=False, dimension=D) self.norm3_tr = get_norm(NORM_TYPE, TR_CHANNELS[3], bn_momentum=bn_momentum, D=D) self.conv2_tr = ME.MinkowskiConvolutionTranspose( in_channels=CHANNELS[2] + TR_CHANNELS[3], out_channels=TR_CHANNELS[2], kernel_size=3, stride=2, dilation=1, has_bias=False, dimension=D) self.norm2_tr = get_norm(NORM_TYPE, TR_CHANNELS[2], bn_momentum=bn_momentum, D=D) self.conv1_tr = ME.MinkowskiConvolution( in_channels=CHANNELS[1] + TR_CHANNELS[2], out_channels=TR_CHANNELS[1], kernel_size=3, stride=1, dilation=1, has_bias=False, dimension=D) self.norm1_tr = get_norm(NORM_TYPE, TR_CHANNELS[1], bn_momentum=bn_momentum, D=D) self.final = ME.MinkowskiConvolution( in_channels=TR_CHANNELS[1], out_channels=out_channels, kernel_size=1, stride=1, dilation=1, has_bias=True, dimension=D)
def __init__(self, input_a_dim, input_b_dim, out_dim, k=3): super().__init__() ''' Transition Up Layer npoint: number of input points nsample: k in kNN, default 3 input_a_dim: feature dimension of the input a(needs upsampling) input_b_dim: feature dimension of the input b (directly concat) out_dim: feature dimension of the TDLayer(fixed as the input_a_dim // 2) + input_b_dim ''' self.k = k self.input_a_dim = input_a_dim self.input_b_dim = input_b_dim self.intermediate_dim = (input_a_dim // 2) + input_b_dim self.out_dim = out_dim self.POINT_TR_LIKE = False self.SUM_FEATURE = True # only used when POINTTR_LIKE is False, somehow have some peoblems # -------- Point TR like ----------- if self.POINT_TR_LIKE: self.linear_a = nn.Linear(input_a_dim, out_dim) self.linear_b = nn.Linear(input_b_dim, out_dim) else: if self.SUM_FEATURE: self.conv_a = nn.Sequential( ME.MinkowskiConvolutionTranspose(in_channels=input_a_dim, out_channels=out_dim, kernel_size=2, stride=2, dimension=3), ME.MinkowskiBatchNorm(out_dim), ME.MinkowskiReLU(), ) self.conv_b = nn.Sequential( ME.MinkowskiConvolution(in_channels=input_b_dim, out_channels=out_dim, kernel_size=1, stride=1, dimension=3), ME.MinkowskiBatchNorm(out_dim), ME.MinkowskiReLU(), ) else: self.conv = ME.MinkowskiConvolutionTranspose( in_channels=input_a_dim, out_channels=input_a_dim // 2, kernel_size=2, stride=2, dimension=3) self.bn = ME.MinkowskiBatchNorm(self.input_a_dim // 2) self.relu = ME.MinkowskiReLU() # ----------------------------------------- self.out_conv = ME.MinkowskiConvolution( in_channels=input_a_dim // 2 + input_b_dim, out_channels=out_dim, kernel_size=3, stride=1, dimension=3) self.out_bn = ME.MinkowskiBatchNorm(self.out_dim) self.out_relu = ME.MinkowskiReLU()
def network_initialization(self, in_channels, out_channels, D): # Output of the first conv concated to conv6 self.inplanes = self.INIT_DIM self.conv0p1s1 = ME.MinkowskiConvolution( in_channels, self.inplanes, kernel_size=5, dimension=D) self.bn0 = ME.MinkowskiBatchNorm(self.inplanes) self.conv1p1s2 = ME.MinkowskiConvolution( self.inplanes, self.inplanes, kernel_size=2, stride=2, dimension=D) self.bn1 = ME.MinkowskiBatchNorm(self.inplanes) self.block1 = self._make_layer(self.BLOCK, self.PLANES[0], self.LAYERS[0], leakiness=self.leakiness) self.conv2p2s2 = ME.MinkowskiConvolution( self.inplanes, self.inplanes, kernel_size=2, stride=2, dimension=D) self.bn2 = ME.MinkowskiBatchNorm(self.inplanes) self.block2 = self._make_layer(self.BLOCK, self.PLANES[1], self.LAYERS[1], leakiness=self.leakiness) self.conv3p4s2 = ME.MinkowskiConvolution( self.inplanes, self.inplanes, kernel_size=2, stride=2, dimension=D) self.bn3 = ME.MinkowskiBatchNorm(self.inplanes) self.block3 = self._make_layer(self.BLOCK, self.PLANES[2], self.LAYERS[2], leakiness=self.leakiness) self.conv4p8s2 = ME.MinkowskiConvolution( self.inplanes, self.inplanes, kernel_size=2, stride=2, dimension=D) self.bn4 = ME.MinkowskiBatchNorm(self.inplanes) self.block4 = self._make_layer(self.BLOCK, self.PLANES[3], self.LAYERS[3], leakiness=self.leakiness) self.convtr4p16s2 = ME.MinkowskiConvolutionTranspose( self.inplanes, self.PLANES[4], kernel_size=2, stride=2, dimension=D) self.bntr4 = ME.MinkowskiBatchNorm(self.PLANES[4]) self.inplanes = self.PLANES[4] if self.skipconnections: self.inplanes += self.PLANES[2] * self.BLOCK.expansion self.block5 = self._make_layer(self.BLOCK, self.PLANES[4], self.LAYERS[4], leakiness=self.leakiness) self.convtr5p8s2 = ME.MinkowskiConvolutionTranspose( self.inplanes, self.PLANES[5], kernel_size=2, stride=2, dimension=D) self.bntr5 = ME.MinkowskiBatchNorm(self.PLANES[5]) self.inplanes = self.PLANES[5] #+ self.PLANES[1] * self.BLOCK.expansion if self.skipconnections: self.inplanes += self.PLANES[1] * self.BLOCK.expansion self.block6 = self._make_layer(self.BLOCK, self.PLANES[5], self.LAYERS[5], leakiness=self.leakiness) self.convtr6p4s2 = ME.MinkowskiConvolutionTranspose( self.inplanes, self.PLANES[6], kernel_size=2, stride=2, dimension=D) self.bntr6 = ME.MinkowskiBatchNorm(self.PLANES[6]) self.inplanes = self.PLANES[6] #+ self.PLANES[0] * self.BLOCK.expansion if self.skipconnections: self.inplanes += self.PLANES[0] * self.BLOCK.expansion self.block7 = self._make_layer(self.BLOCK, self.PLANES[6], self.LAYERS[6], leakiness=self.leakiness) self.convtr7p2s2 = ME.MinkowskiConvolutionTranspose( self.inplanes, self.PLANES[7], kernel_size=2, stride=2, dimension=D) self.bntr7 = ME.MinkowskiBatchNorm(self.PLANES[7]) self.inplanes = self.PLANES[7] #+ self.INIT_DIM if self.skipconnections: self.inplanes += self.INIT_DIM self.block8 = self._make_layer(self.BLOCK, self.PLANES[7], self.LAYERS[7], leakiness=self.leakiness) self.final = ME.MinkowskiConvolution( self.PLANES[7], out_channels, kernel_size=1, has_bias=True, dimension=D) self.relu = ME.MinkowskiLeakyReLU(negative_slope=self.leakiness) self.dropout = ME.MinkowskiDropout(self.dropout_p)
def network_initialization(self, in_channels, out_channels, D): assert len(self.layers) == len(self.planes) assert len(self.planes) == self.num_bottom_up self.convs = nn.ModuleList( ) # Bottom-up convolutional blocks with stride=2 self.bn = nn.ModuleList() # Bottom-up BatchNorms self.blocks = nn.ModuleList() # Bottom-up blocks self.tconvs = nn.ModuleList() # Top-down tranposed convolutions self.conv1x1 = nn.ModuleList( ) # 1x1 convolutions in lateral connections # The first convolution is special case, with kernel size = 5 self.inplanes = self.planes[0] self.conv0 = ME.MinkowskiConvolution( in_channels, self.inplanes, kernel_size=self.conv0_kernel_size, dimension=D) self.bn0 = ME.MinkowskiBatchNorm(self.inplanes) for plane, layer in zip(self.planes, self.layers): self.convs.append( ME.MinkowskiConvolution(self.inplanes, self.inplanes, kernel_size=2, stride=2, dimension=D)) self.bn.append(ME.MinkowskiBatchNorm(self.inplanes)) self.blocks.append(self._make_layer(self.block, plane, layer)) # Lateral connections for i in range(self.num_top_down): self.conv1x1.append( ME.MinkowskiConvolution(self.planes[-1 - i], self.lateral_dim, kernel_size=1, stride=1, dimension=D)) self.tconvs.append( ME.MinkowskiConvolutionTranspose(self.lateral_dim, self.lateral_dim, kernel_size=2, stride=2, dimension=D)) # There's one more lateral connection than top-down TConv blocks if self.num_top_down < self.num_bottom_up: # Lateral connection from Conv block 1 or above self.conv1x1.append( ME.MinkowskiConvolution(self.planes[-1 - self.num_top_down], self.lateral_dim, kernel_size=1, stride=1, dimension=D)) else: # Lateral connection from Con0 block self.conv1x1.append( ME.MinkowskiConvolution(self.planes[0], self.lateral_dim, kernel_size=1, stride=1, dimension=D)) self.relu = ME.MinkowskiReLU(inplace=True)
def __init__(self, cfg, name='fpn'): super(FPN, self).__init__(cfg) model_cfg = cfg['modules'][name] # UResNet Configurations self.reps = model_cfg.get('reps', 2) self.depth = model_cfg.get('depth', 5) self.num_filters = model_cfg.get('num_filters', 16) self.nPlanes = [i * self.num_filters for i in range(1, self.depth + 1)] # self.kernel_size = cfg.get('kernel_size', 3) # self.downsample = cfg.get(downsample, 2) self.input_kernel = model_cfg.get('input_kernel', 3) # Initialize Input Layer self.input_layer = ME.MinkowskiConvolution( in_channels=self.num_input, out_channels=self.num_filters, kernel_size=self.input_kernel, stride=1, dimension=self.D) # Initialize Encoder self.encoding_conv = [] self.encoding_block = [] for i, F in enumerate(self.nPlanes): m = [] for _ in range(self.reps): m.append( ResNetBlock(F, F, dimension=self.D, leakiness=self.leakiness)) m = nn.Sequential(*m) self.encoding_block.append(m) m = [] if i < self.depth - 1: m.append(ME.MinkowskiBatchNorm(F)) m.append(MinkowskiLeakyReLU(negative_slope=self.leakiness)) m.append( ME.MinkowskiConvolution(in_channels=self.nPlanes[i], out_channels=self.nPlanes[i + 1], kernel_size=2, stride=2, dimension=self.D)) m = nn.Sequential(*m) self.encoding_conv.append(m) self.encoding_conv = nn.Sequential(*self.encoding_conv) self.encoding_block = nn.Sequential(*self.encoding_block) # Lateral Connections self.lateral = [] for i, F in enumerate(self.nPlanes[-2::-1]): self.lateral.append(ME.MinkowskiLinear(F, F)) self.lateral = nn.Sequential(*self.lateral) # Initialize Decoder self.decoding_block = [] self.decoding_conv = [] for i in range(self.depth - 2, -1, -1): m = [] m.append(ME.MinkowskiBatchNorm(self.nPlanes[i + 1])) m.append(MinkowskiLeakyReLU(negative_slope=self.leakiness)) m.append( ME.MinkowskiConvolutionTranspose(in_channels=self.nPlanes[i + 1], out_channels=self.nPlanes[i], kernel_size=2, stride=2, dimension=self.D)) m = nn.Sequential(*m) self.decoding_conv.append(m) m = [] for j in range(self.reps): m.append( ResNetBlock(self.nPlanes[i], self.nPlanes[i], dimension=self.D, leakiness=self.leakiness)) m = nn.Sequential(*m) self.decoding_block.append(m) self.decoding_block = nn.Sequential(*self.decoding_block) self.decoding_conv = nn.Sequential(*self.decoding_conv) print('Total Number of Trainable Parameters = {}'.format( sum(p.numel() for p in self.parameters() if p.requires_grad)))
def __init__(self,config,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 bn_momentum = config.bn_momentum self.normalize_feature = config.normalize_feature self.voxel_size = config.voxel_size self.conv1 = ME.MinkowskiConvolution( in_channels=config.in_feats_dim, out_channels=CHANNELS[1], kernel_size=config.conv1_kernel_size, stride=1, dilation=1, bias=False, dimension=D) self.norm1 = get_norm(NORM_TYPE, CHANNELS[1], bn_momentum=bn_momentum, D=D) self.block1 = get_block( BLOCK_NORM_TYPE, CHANNELS[1], CHANNELS[1], bn_momentum=bn_momentum, D=D) self.conv2 = ME.MinkowskiConvolution( in_channels=CHANNELS[1], out_channels=CHANNELS[2], kernel_size=3, stride=2, dilation=1, bias=False, dimension=D) self.norm2 = get_norm(NORM_TYPE, CHANNELS[2], bn_momentum=bn_momentum, D=D) self.block2 = get_block( BLOCK_NORM_TYPE, CHANNELS[2], CHANNELS[2], bn_momentum=bn_momentum, D=D) self.conv3 = ME.MinkowskiConvolution( in_channels=CHANNELS[2], out_channels=CHANNELS[3], kernel_size=3, stride=2, dilation=1, bias=False, dimension=D) self.norm3 = get_norm(NORM_TYPE, CHANNELS[3], bn_momentum=bn_momentum, D=D) self.block3 = get_block( BLOCK_NORM_TYPE, CHANNELS[3], CHANNELS[3], bn_momentum=bn_momentum, D=D) self.conv4 = ME.MinkowskiConvolution( in_channels=CHANNELS[3], out_channels=CHANNELS[4], kernel_size=3, stride=2, dilation=1, bias=False, dimension=D) self.norm4 = get_norm(NORM_TYPE, CHANNELS[4], bn_momentum=bn_momentum, D=D) self.block4 = get_block( BLOCK_NORM_TYPE, CHANNELS[4], CHANNELS[4], bn_momentum=bn_momentum, D=D) # adapt input tensor here self.conv4_tr = ME.MinkowskiConvolutionTranspose( in_channels = config.gnn_feats_dim + 2, out_channels=TR_CHANNELS[4], kernel_size=3, stride=2, dilation=1, bias=False, dimension=D) self.norm4_tr = get_norm(NORM_TYPE, TR_CHANNELS[4], bn_momentum=bn_momentum, D=D) self.block4_tr = get_block( BLOCK_NORM_TYPE, TR_CHANNELS[4], TR_CHANNELS[4], bn_momentum=bn_momentum, D=D) self.conv3_tr = ME.MinkowskiConvolutionTranspose( in_channels=CHANNELS[3] + TR_CHANNELS[4], out_channels=TR_CHANNELS[3], kernel_size=3, stride=2, dilation=1, bias=False, dimension=D) self.norm3_tr = get_norm(NORM_TYPE, TR_CHANNELS[3], bn_momentum=bn_momentum, D=D) self.block3_tr = get_block( BLOCK_NORM_TYPE, TR_CHANNELS[3], TR_CHANNELS[3], bn_momentum=bn_momentum, D=D) self.conv2_tr = ME.MinkowskiConvolutionTranspose( in_channels=CHANNELS[2] + TR_CHANNELS[3], out_channels=TR_CHANNELS[2], kernel_size=3, stride=2, dilation=1, bias=False, dimension=D) self.norm2_tr = get_norm(NORM_TYPE, TR_CHANNELS[2], bn_momentum=bn_momentum, D=D) self.block2_tr = get_block( BLOCK_NORM_TYPE, TR_CHANNELS[2], TR_CHANNELS[2], bn_momentum=bn_momentum, D=D) self.conv1_tr = ME.MinkowskiConvolution( in_channels=CHANNELS[1] + TR_CHANNELS[2], out_channels=TR_CHANNELS[1], kernel_size=1, stride=1, dilation=1, bias=False, dimension=D) self.final = ME.MinkowskiConvolution( in_channels=TR_CHANNELS[1], out_channels=config.out_feats_dim + 2, kernel_size=1, stride=1, dilation=1, bias=True, dimension=D) ############# # Overlap attention module self.epsilon = torch.nn.Parameter(torch.tensor(-5.0)) self.bottle = nn.Conv1d(CHANNELS[4], config.gnn_feats_dim,kernel_size=1,bias=True) self.gnn = GCN(config.num_head,config.gnn_feats_dim, config.dgcnn_k, config.nets) self.proj_gnn = nn.Conv1d(config.gnn_feats_dim,config.gnn_feats_dim,kernel_size=1, bias=True) self.proj_score = nn.Conv1d(config.gnn_feats_dim,1,kernel_size=1,bias=True)
def __init__(self, in_channels=3, out_channels=32, bn_momentum=0.1, D=3, config=None): ME.MinkowskiNetwork.__init__(self, D) NORM_TYPE = self.NORM_TYPE bn_momentum = config.bn_momentum CHANNELS = self.CHANNELS TR_CHANNELS = self.TR_CHANNELS self.normalize_feature = config.normalize_feature self.conv1 = ME.MinkowskiConvolution( in_channels=in_channels, out_channels=CHANNELS[1], kernel_size=config.conv1_kernel_size, stride=1, dilation=1, has_bias=False, dimension=D) self.norm1 = get_norm(NORM_TYPE, CHANNELS[1], bn_momentum=bn_momentum, D=D) self.conv2 = ME.MinkowskiConvolution( in_channels=CHANNELS[1], out_channels=CHANNELS[2], kernel_size=3, stride=2, dilation=1, has_bias=False, dimension=D) self.norm2 = get_norm(NORM_TYPE, CHANNELS[2], bn_momentum=bn_momentum, D=D) self.conv3 = ME.MinkowskiConvolution( in_channels=CHANNELS[2], out_channels=CHANNELS[3], kernel_size=3, stride=2, dilation=1, has_bias=False, dimension=D) self.norm3 = get_norm(NORM_TYPE, CHANNELS[3], bn_momentum=bn_momentum, D=D) self.conv4 = ME.MinkowskiConvolution( in_channels=CHANNELS[3], out_channels=CHANNELS[4], kernel_size=3, stride=2, dilation=1, has_bias=False, dimension=D) self.norm4 = get_norm(NORM_TYPE, CHANNELS[4], bn_momentum=bn_momentum, D=D) self.conv5 = ME.MinkowskiConvolution( in_channels=CHANNELS[4], out_channels=CHANNELS[5], kernel_size=3, stride=2, dilation=1, has_bias=False, dimension=D) self.norm5 = get_norm(NORM_TYPE, CHANNELS[5], bn_momentum=bn_momentum, D=D) self.conv5_tr = ME.MinkowskiConvolutionTranspose( in_channels=CHANNELS[5], out_channels=TR_CHANNELS[5], kernel_size=3, stride=2, dilation=1, has_bias=False, dimension=D) self.norm5_tr = get_norm(NORM_TYPE, TR_CHANNELS[5], bn_momentum=bn_momentum, D=D) self.conv4_tr = ME.MinkowskiConvolutionTranspose( in_channels=CHANNELS[4] + TR_CHANNELS[5], out_channels=TR_CHANNELS[4], kernel_size=3, stride=2, dilation=1, has_bias=False, dimension=D) self.norm4_tr = get_norm(NORM_TYPE, TR_CHANNELS[4], bn_momentum=bn_momentum, D=D) self.conv3_tr = ME.MinkowskiConvolutionTranspose( in_channels=CHANNELS[3] + TR_CHANNELS[4], out_channels=TR_CHANNELS[3], kernel_size=3, stride=2, dilation=1, has_bias=False, dimension=D) self.norm3_tr = get_norm(NORM_TYPE, TR_CHANNELS[3], bn_momentum=bn_momentum, D=D) self.conv2_tr = ME.MinkowskiConvolutionTranspose( in_channels=CHANNELS[2] + TR_CHANNELS[3], out_channels=TR_CHANNELS[2], kernel_size=3, stride=2, dilation=1, has_bias=False, dimension=D) self.norm2_tr = get_norm(NORM_TYPE, TR_CHANNELS[2], bn_momentum=bn_momentum, D=D) self.conv1_tr = ME.MinkowskiConvolution( in_channels=CHANNELS[1] + TR_CHANNELS[2], out_channels=TR_CHANNELS[1], kernel_size=1, stride=1, dilation=1, has_bias=True, dimension=D)
def __init__(self, cfg, name='uresnext'): super(UResNeXt, self).__init__(cfg) self.model_cfg = cfg['modules'][name] # Configurations self.reps = self.model_cfg.get('reps', 2) self.depth = self.model_cfg.get('depth', 5) self.num_filters = self.model_cfg.get('num_filters', 16) self.cardinality = self.model_cfg.get('cardinality', 4) assert (self.num_filters % self.cardinality == 0) self.nPlanes = [i * self.num_filters for i in range(1, self.depth + 1)] self.input_kernel = self.model_cfg.get('input_kernel', 3) # Initialize Input Layer self.input_layer = AtrousIIBlock(self.num_input, self.num_filters, leakiness=self.leakiness) # Initialize Encoder self.encoding_conv = [] self.encoding_block = [] for i, F in enumerate(self.nPlanes): m = [] for _ in range(self.reps): m.append( ResNeXtBlock(F, F, dimension=self.D, leakiness=self.leakiness, cardinality=4, dilations=[1, 1, 3, 9])) m = nn.Sequential(*m) self.encoding_block.append(m) m = [] if i < self.depth - 1: m.append(ME.MinkowskiBatchNorm(F)) m.append(MinkowskiLeakyReLU(negative_slope=self.leakiness)) m.append( ME.MinkowskiConvolution(in_channels=self.nPlanes[i], out_channels=self.nPlanes[i + 1], kernel_size=2, stride=2, dimension=self.D)) m = nn.Sequential(*m) self.encoding_conv.append(m) self.encoding_conv = nn.Sequential(*self.encoding_conv) self.encoding_block = nn.Sequential(*self.encoding_block) # Initialize Decoder self.decoding_block = [] self.decoding_conv = [] for i in range(self.depth - 2, -1, -1): m = [] m.append(ME.MinkowskiBatchNorm(self.nPlanes[i + 1])) m.append(MinkowskiLeakyReLU(negative_slope=self.leakiness)) m.append( ME.MinkowskiConvolutionTranspose(in_channels=self.nPlanes[i + 1], out_channels=self.nPlanes[i], kernel_size=2, stride=2, dimension=self.D)) m = nn.Sequential(*m) self.decoding_conv.append(m) m = [] for j in range(self.reps): m.append( ResNeXtBlock(self.nPlanes[i] * (2 if j == 0 else 1), self.nPlanes[i], dimension=self.D, cardinality=4, dilations=[1, 1, 3, 9], leakiness=self.leakiness)) m = nn.Sequential(*m) self.decoding_block.append(m) self.decoding_block = nn.Sequential(*self.decoding_block) self.decoding_conv = nn.Sequential(*self.decoding_conv) print('Total Number of Trainable Parameters = {}'.format( sum(p.numel() for p in self.parameters() if p.requires_grad)))