def network_initialization(self, in_channels, out_channels, D): self.inplanes = self.INIT_DIM self.conv1 = ME.MinkowskiConvolution( in_channels, self.inplanes, kernel_size=5, stride=1, dimension=D) self.bn1 = ME.MinkowskiBatchNorm(self.inplanes) self.relu = ME.MinkowskiReLU(inplace=True) self.pool = ME.MinkowskiSumPooling(kernel_size=2, stride=2, dimension=D) self.layer1 = self._make_layer( self.BLOCK, self.PLANES[0], self.LAYERS[0], stride=2) self.layer2 = self._make_layer( self.BLOCK, self.PLANES[1], self.LAYERS[1], stride=2) self.layer3 = self._make_layer( self.BLOCK, self.PLANES[2], self.LAYERS[2], stride=2) self.layer4 = self._make_layer( self.BLOCK, self.PLANES[3], self.LAYERS[3], stride=2) self.glob_avg = ME.MinkowskiGlobalPooling(dimension=D) self.classification_block = nn.Sequential( ME.MinkowskiLinear(self.inplanes, self.inplanes, bias=False), ME.MinkowskiBatchNorm(self.inplanes), ME.MinkowskiReLU(), ME.MinkowskiLinear(self.inplanes, self.inplanes, bias=False), ME.MinkowskiBatchNorm(self.inplanes)) self.final = ME.MinkowskiLinear(self.inplanes, out_channels, bias=True)
def sum_pool(kernel_size, stride=1, dilation=1, 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, stride, dilation, region_type=region_type, axis_types=axis_types, dimension=D) return ME.MinkowskiSumPooling( kernel_size=kernel_size, stride=stride, dilation=dilation, kernel_generator=kernel_generator, dimension=D)
def check_data(model, train_data_loader, val_data_loader, config): data_iter = train_data_loader.__iter__() import ipdb ipdb.set_trace() sample_size = 1 strides = [2, 4, 8, 16] # strides = [2] for stride in strides: coordinate_map_key = [stride] * 3 pool0 = ME.MinkowskiSumPooling(kernel_size=stride, stride=stride, dilation=1, dimension=3) all_neis_scenes = [] for _ in range(sample_size): coords, input, target, _, _ = data_iter.next() assert coords[:, 0].max() == 0 # assert bs=1 x = ME.SparseTensor(input, coords, device='cuda') x = pool0(x) d = {} neis_d = x.coordinate_manager.get_kernel_map( x.coordinate_map_key, x.coordinate_map_key, kernel_size=3, stride=1, dilation=1, ) # d['all_c'] = x.C[:,1:] N = x.C.shape[0] k = 27 all_neis = [] for k_ in range(k): if not k_ in neis_d.keys(): continue neis_ = torch.gather(x.C[:, 1:].float(), dim=0, index=neis_d[k_][0].reshape(-1, 1).repeat( 1, 3).long()) neis = torch.zeros(N, 3, device=x.F.device) neis.scatter_(dim=0, index=neis_d[k_][1].reshape(-1, 1).repeat(1, 3).long(), src=neis_) neis = (neis.sum(-1) > 0).int() all_neis.append(neis) all_neis = torch.stack(all_neis) all_neis_scenes.append(all_neis) # d['neis'] = all_neis d['sparse_mask'] = torch.cat(all_neis_scenes, dim=-1) print('Stride:{} Shape:'.format(stride), d['sparse_mask'].shape) name = "sparse_mask_s{}_nuscenes.pth".format(stride) # name = 'test-kitti' torch.save( d, '/home/zhaotianchen/project/point-transformer/SpatioTemporalSegmentation-ScanNet/plot/final/{}' .format(name)) # import ipdb; ipdb.set_trace() import ipdb ipdb.set_trace()
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 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=ME.RegionType.HYPERCUBE, dimension=D) self.norm1 = get_norm(NORM_TYPE, CHANNELS[1], bn_momentum=bn_momentum, dimension=D) self.block1 = get_block( BLOCK_NORM_TYPE, CHANNELS[1], CHANNELS[1], bn_momentum=bn_momentum, region_type=REGION_TYPE, dimension=D) self.pool2 = ME.MinkowskiSumPooling(kernel_size=2, stride=2, dimension=D) self.conv2 = conv( in_channels=CHANNELS[1], out_channels=CHANNELS[2], kernel_size=3, stride=1, dilation=1, has_bias=False, region_type=REGION_TYPE, dimension=D) self.norm2 = get_norm(NORM_TYPE, CHANNELS[2], bn_momentum=bn_momentum, dimension=D) self.block2 = get_block( BLOCK_NORM_TYPE, CHANNELS[2], CHANNELS[2], bn_momentum=bn_momentum, region_type=REGION_TYPE, dimension=D) self.pool3 = ME.MinkowskiSumPooling(kernel_size=2, stride=2, dimension=D) self.conv3 = conv( in_channels=CHANNELS[2], out_channels=CHANNELS[3], kernel_size=3, stride=1, dilation=1, has_bias=False, region_type=REGION_TYPE, dimension=D) self.norm3 = get_norm(NORM_TYPE, CHANNELS[3], bn_momentum=bn_momentum, dimension=D) self.block3 = get_block( BLOCK_NORM_TYPE, CHANNELS[3], CHANNELS[3], bn_momentum=bn_momentum, region_type=REGION_TYPE, dimension=D) self.pool4 = ME.MinkowskiSumPooling(kernel_size=2, stride=2, dimension=D) self.conv4 = conv( in_channels=CHANNELS[3], out_channels=CHANNELS[4], kernel_size=3, stride=1, dilation=1, has_bias=False, region_type=ME.RegionType.HYPERCUBE, dimension=D) self.norm4 = get_norm(NORM_TYPE, CHANNELS[4], bn_momentum=bn_momentum, dimension=D) self.block4 = get_block( BLOCK_NORM_TYPE, CHANNELS[4], CHANNELS[4], bn_momentum=bn_momentum, region_type=ME.RegionType.HYPERCUBE, dimension=D) self.conv4_tr = conv_tr( in_channels=CHANNELS[4], out_channels=TR_CHANNELS[4], kernel_size=3, stride=2, dilation=1, has_bias=False, region_type=ME.RegionType.HYPERCUBE, dimension=D) self.norm4_tr = get_norm( NORM_TYPE, TR_CHANNELS[4], bn_momentum=bn_momentum, dimension=D) self.block4_tr = get_block( BLOCK_NORM_TYPE, TR_CHANNELS[4], TR_CHANNELS[4], bn_momentum=bn_momentum, region_type=REGION_TYPE, dimension=D) self.conv3_tr = conv_tr( in_channels=CHANNELS[3] + TR_CHANNELS[4], out_channels=TR_CHANNELS[3], kernel_size=3, stride=2, dilation=1, has_bias=False, region_type=REGION_TYPE, dimension=D) self.norm3_tr = get_norm( NORM_TYPE, TR_CHANNELS[3], bn_momentum=bn_momentum, dimension=D) self.block3_tr = get_block( BLOCK_NORM_TYPE, TR_CHANNELS[3], TR_CHANNELS[3], bn_momentum=bn_momentum, region_type=REGION_TYPE, dimension=D) self.conv2_tr = conv_tr( in_channels=CHANNELS[2] + TR_CHANNELS[3], out_channels=TR_CHANNELS[2], kernel_size=3, stride=2, 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 = get_block( BLOCK_NORM_TYPE, TR_CHANNELS[2], TR_CHANNELS[2], bn_momentum=bn_momentum, region_type=REGION_TYPE, dimension=D) self.conv1_tr = conv( 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, 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)