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)
