def __init__(self, input_dim, p=3, eps=1e-6): super(GeM, self).__init__() self.input_dim = input_dim # Same output number of channels as input number of channels self.output_dim = self.input_dim self.p = nn.Parameter(torch.ones(1) * p) self.eps = eps self.f = ME.MinkowskiGlobalAvgPooling()
def network_initialization( self, in_channel, out_channel, channels, embedding_channel, kernel_size, D=3, ): self.mlp1 = self.get_mlp_block(in_channel, channels[0]) self.conv1 = self.get_conv_block( channels[0], channels[1], kernel_size=kernel_size, stride=1, ) self.conv2 = self.get_conv_block( channels[1], channels[2], kernel_size=kernel_size, stride=2, ) self.conv3 = self.get_conv_block( channels[2], channels[3], kernel_size=kernel_size, stride=2, ) self.conv4 = self.get_conv_block( channels[3], channels[4], kernel_size=kernel_size, stride=2, ) self.conv5 = nn.Sequential( self.get_conv_block( channels[1] + channels[2] + channels[3] + channels[4], embedding_channel // 4, kernel_size=3, stride=2, ), self.get_conv_block( embedding_channel // 4, embedding_channel // 2, kernel_size=3, stride=2, ), self.get_conv_block( embedding_channel // 2, embedding_channel, kernel_size=3, stride=2, ), ) self.pool = ME.MinkowskiMaxPooling(kernel_size=3, stride=2, dimension=D) self.global_max_pool = ME.MinkowskiGlobalMaxPooling() self.global_avg_pool = ME.MinkowskiGlobalAvgPooling() self.final = nn.Sequential( self.get_mlp_block(embedding_channel * 2, 512), ME.MinkowskiDropout(), self.get_mlp_block(512, 512), ME.MinkowskiLinear(512, out_channel, bias=True), )
def __init__(self, p=3, eps=1e-6): super(GeM, self).__init__() self.p = nn.Parameter(torch.ones(1) * p) self.eps = eps self.f = ME.MinkowskiGlobalAvgPooling()
def __init__(self): super().__init__() self.f = ME.MinkowskiGlobalAvgPooling()
def __init__(self, input_dim): super().__init__() self.input_dim = input_dim # Same output number of channels as input number of channels self.output_dim = self.input_dim self.f = ME.MinkowskiGlobalAvgPooling()
# Normalize features and create a sparse tensor return features, coordinates if __name__ == '__main__': config = parser.parse_args() device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # Define a model and load the weights feats = [3, 8, 16, 32, 64, 128] features, coordinates = generate_input_sparse_tensor( config.file_name, voxel_size=config.voxel_size, batch_size=config.batch_size) pool = ME.MinkowskiGlobalAvgPooling() # Measure time print('Forward') for feat in feats: timer = Timer() features = torch.rand(len(coordinates), feat).to(device) # Feed-forward pass and get the prediction for i in range(20): sinput = ME.SparseTensor(features, coordinates=coordinates, device=device) timer.tic() soutput = pool(sinput)
def network_initialization(self, in_channels, out_channels, D): # Setup net_metadata dilations = self.DILATIONS bn_momentum = 0.02 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 conv1_kernel_size = 3 self.inplanes = self.INIT_DIM self.conv0p1s1 = conv(in_channels, self.inplanes, kernel_size=space_n_time_m(conv1_kernel_size, 1), stride=1, dilation=1, conv_type=self.NON_BLOCK_CONV_TYPE, D=D) self.bn0 = get_norm(self.NORM_TYPE, self.inplanes, D, bn_momentum=bn_momentum) self.conv1p1s2 = 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.bn1 = get_norm(self.NORM_TYPE, self.inplanes, 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.conv2p2s2 = 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.conv3p4s2 = 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.conv4p8s2 = 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.convtr4p16s2 = 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.convtr5p8s2 = 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.convtr6p4s2 = 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.inplanes = self.PLANES[6] + self.PLANES[0] * self.BLOCK.expansion self.block7 = self._make_layer(self.BLOCK, self.PLANES[6], self.LAYERS[6], dilation=dilations[6], norm_type=self.NORM_TYPE, bn_momentum=bn_momentum) self.convtr7p2s2 = conv_tr(self.inplanes, self.PLANES[7], 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.bntr7 = get_norm(self.NORM_TYPE, self.PLANES[7], D, bn_momentum=bn_momentum) self.inplanes = self.PLANES[7] + self.INIT_DIM self.block8 = self._make_layer(self.BLOCK, self.PLANES[7], self.LAYERS[7], dilation=dilations[7], norm_type=self.NORM_TYPE, bn_momentum=bn_momentum) self.all_feat_names = [ "en0", "en1", "en2", "en3", "en4", "plane4", "plane5", "plane6", "plane7", ] self.relu = MinkowskiReLU(inplace=True) self.maxpool = ME.MinkowskiGlobalMaxPooling() self.avgpool = ME.MinkowskiGlobalAvgPooling() if self.use_mlp: self.head = SMLP(self.mlp_dim)
def __init__(self, in_channel, out_channel, num_class, embedding_channel=1024, dimension=3): ME.MinkowskiNetwork.__init__(self, dimension) # The normal channel for Modelnet is 3, for scannet is 6, for scanobjnn is 0 normal_channel = 3 # in_channel = normal_channel+3 # normal ch + xyz self.normal_channel = normal_channel self.input_mlp = nn.Sequential( ME.MinkowskiConvolution(in_channel, 32, kernel_size=1, dimension=3), ME.MinkowskiBatchNorm(32), ME.MinkowskiReLU(), ME.MinkowskiConvolution(32, 32, kernel_size=1, dimension=3), ME.MinkowskiBatchNorm(32)) self.in_dims = [32, 64, 128, 256] self.out_dims = [64, 128, 256, 512] self.neighbor_ks = [16, 32, 64, 16, 16] # self.neighbor_ks = [8, 8, 8, 8, 8] self.PTBlock0 = PTBlock(in_dim=self.in_dims[0], n_sample=self.neighbor_ks[0]) self.TDLayer1 = TDLayer(input_dim=self.in_dims[0], out_dim=self.out_dims[0]) self.PTBlock1 = PTBlock(in_dim=self.out_dims[0], n_sample=self.neighbor_ks[1]) self.TDLayer2 = TDLayer(input_dim=self.in_dims[1], out_dim=self.out_dims[1]) self.PTBlock2 = PTBlock(in_dim=self.out_dims[1], n_sample=self.neighbor_ks[1]) self.TDLayer3 = TDLayer(input_dim=self.in_dims[2], out_dim=self.out_dims[2]) self.PTBlock3 = PTBlock(in_dim=self.out_dims[2], n_sample=self.neighbor_ks[2]) self.TDLayer4 = TDLayer(input_dim=self.in_dims[3], out_dim=self.out_dims[3]) self.PTBlock4 = PTBlock(in_dim=self.out_dims[3], n_sample=self.neighbor_ks[4]) self.middle_linear = ME.MinkowskiConvolution(self.out_dims[3], self.out_dims[3], kernel_size=1, dimension=3) self.PTBlock_middle = PTBlock(in_dim=self.out_dims[3], n_sample=self.neighbor_ks[4]) self.TULayer5 = TULayer(input_dim=self.out_dims[3], out_dim=self.in_dims[3]) self.PTBlock5 = PTBlock(in_dim=self.in_dims[3], n_sample=self.neighbor_ks[4]) self.TULayer6 = TULayer(input_dim=self.out_dims[2], out_dim=self.in_dims[2]) self.PTBlock6 = PTBlock(in_dim=self.in_dims[2], n_sample=self.neighbor_ks[3]) self.TULayer7 = TULayer(input_dim=self.out_dims[1], out_dim=self.in_dims[1]) self.PTBlock7 = PTBlock(in_dim=self.in_dims[1], n_sample=self.neighbor_ks[2]) self.TULayer8 = TULayer(input_dim=self.out_dims[0], out_dim=self.in_dims[0]) self.PTBlock8 = PTBlock(in_dim=self.in_dims[0], n_sample=self.neighbor_ks[1]) self.fc = nn.Sequential( # ME.MinkowskiLinear(32, 32), ME.MinkowskiLinear(self.out_dims[3], 32), ME.MinkowskiDropout(0.4), ME.MinkowskiLinear(32, num_class)) self.global_avg_pool = ME.MinkowskiGlobalAvgPooling()