def __init__(self):
super().__init__()
self.net = spconv.SparseSequential(
spconv.SubMConv3d(3, 8, 3, indice_key="subm1", padding=1, use_hash=False),
nn.BatchNorm1d(8),
nn.ReLU(),
spconv.SparseConv3d(8, 16, 3, stride=2, padding=1, use_hash=False),
nn.BatchNorm1d(16),
nn.ReLU(),
spconv.SubMConv3d(16, 16, 3, indice_key="subm2", padding=1, use_hash=False),
nn.BatchNorm1d(16),
nn.ReLU(),
spconv.SparseConv3d(16, 32, 3, stride=2, padding=1, use_hash=False),
nn.BatchNorm1d(32),
nn.ReLU(),
spconv.SubMConv3d(32, 32, 3, indice_key="subm3", padding=1, use_hash=False),
nn.BatchNorm1d(32),
nn.ReLU(),
spconv.SparseConv3d(32, 64, 3, stride=2, padding=1, use_hash=False),
nn.BatchNorm1d(64),
nn.ReLU(),
spconv.SubMConv3d(64, 64, 3, indice_key="subm4", padding=1, use_hash=False),
nn.BatchNorm1d(64),
nn.ReLU(),
spconv.ToDense() # [64, 2, 8, 8]
)
self.linear = nn.Linear(64 * 2 * 8 * 8, 4)
def __init__(self, num_layers, ndim, shape, in_channels, out_channels, kernel_size,
stride, padding, dilation):
super().__init__()
layers = [spconv.SparseConv3d(
in_channels,
out_channels,
kernel_size,
stride,
padding=padding,
dilation=dilation,
bias=False,
use_hash=True)]
for i in range(1, num_layers):
layers.append(spconv.SparseConv3d(
out_channels,
out_channels,
kernel_size,
stride,
padding=padding,
dilation=dilation,
bias=False))
self.net = spconv.SparseSequential(
*layers,
)
# self.grid = torch.full([3, *shape], -1, dtype=torch.int32).cuda()
self.grid = None
self.shape = shape
def __init__(self, in_filters, out_filters, dropout_rate, kernel_size=(3, 3, 3), stride=1,
pooling=True, drop_out=True, height_pooling=False, indice_key=None):
super(ResBlock, self).__init__()
self.pooling = pooling
self.drop_out = drop_out
self.conv1 = conv3x1(in_filters, out_filters, indice_key=indice_key + "bef")
self.act1 = nn.LeakyReLU()
self.bn0 = nn.BatchNorm1d(out_filters)
self.conv1_2 = conv1x3(out_filters, out_filters, indice_key=indice_key + "bef")
self.act1_2 = nn.LeakyReLU()
self.bn0_2 = nn.BatchNorm1d(out_filters)
self.conv2 = conv1x3(in_filters, out_filters, indice_key=indice_key + "bef")
self.act2 = nn.LeakyReLU()
self.bn1 = nn.BatchNorm1d(out_filters)
self.conv3 = conv3x1(out_filters, out_filters, indice_key=indice_key + "bef")
self.act3 = nn.LeakyReLU()
self.bn2 = nn.BatchNorm1d(out_filters)
if pooling:
if height_pooling:
self.pool = spconv.SparseConv3d(out_filters, out_filters, kernel_size=3, stride=2,
padding=1, indice_key=indice_key, bias=False)
else:
self.pool = spconv.SparseConv3d(out_filters, out_filters, kernel_size=3, stride=(2, 2, 1),
padding=1, indice_key=indice_key, bias=False)
self.weight_initialization()
def __init__(self, num_input_features):
super(VxNet, self).__init__()
self.conv0 = double_conv(num_input_features, 16, 'subm0')
self.down0 = stride_conv(16, 32, 'down0')
self.conv1 = double_conv(32, 32, 'subm1')
self.down1 = stride_conv(32, 64, 'down1')
self.conv2 = triple_conv(64, 64, 'subm2')
self.down2 = stride_conv(64, 64, 'down2')
self.conv3 = triple_conv(64, 64, 'subm3') # middle line
self.extra_conv = spconv.SparseSequential(
spconv.SparseConv3d(64, 64, (2, 1, 1), (2, 1, 1),
bias=False), # shape no change
nn.BatchNorm1d(64, eps=1e-3, momentum=0.01),
nn.ReLU())
self.conv_output = spconv.SparseSequential(
spconv.SparseConv3d(192, 64, (1, 1, 1), (1, 1, 1),
bias=False), # shape no change
nn.BatchNorm1d(64, eps=1e-3, momentum=0.01),
nn.ReLU())
self.point_fc = nn.Linear(160, 64, bias=False)
self.point_cls = nn.Linear(64, 1, bias=False)
self.point_reg = nn.Linear(64, 3, bias=False)
def __init__(self,
output_shape,
use_norm=True,
num_input_features=128,
num_filters_down1=[64],
num_filters_down2=[64, 64],
name='SparseMiddleExtractor'):
super(SparseMiddleExtractor, self).__init__()
self.name = name
if use_norm:
BatchNorm1d = change_default_args(eps=1e-3,
momentum=0.01)(nn.BatchNorm1d)
Linear = change_default_args(bias=False)(nn.Linear)
else:
BatchNorm1d = Empty
Linear = change_default_args(bias=True)(nn.Linear)
sparse_shape = np.array(output_shape[1:4]) + [1, 0, 0]
# sparse_shape[0] = 11
print(sparse_shape)
self.sparse_shape = sparse_shape
self.scn_input = scn.InputLayer(3, sparse_shape.tolist())
self.voxel_output_shape = output_shape
middle_layers = []
num_filters = [num_input_features] + num_filters_down1
# num_filters = [64] + num_filters_down1
filters_pairs_d1 = [[num_filters[i], num_filters[i + 1]]
for i in range(len(num_filters) - 1)]
for i, o in filters_pairs_d1:
middle_layers.append(
spconv.SubMConv3d(i, o, 3, bias=False, indice_key="subm0"))
middle_layers.append(BatchNorm1d(o))
middle_layers.append(nn.ReLU())
middle_layers.append(
spconv.SparseConv3d(num_filters[-1],
num_filters[-1], (3, 1, 1), (2, 1, 1),
bias=False))
middle_layers.append(BatchNorm1d(num_filters[-1]))
middle_layers.append(nn.ReLU())
# assert len(num_filters_down2) > 0
if len(num_filters_down1) == 0:
num_filters = [num_filters[-1]] + num_filters_down2
else:
num_filters = [num_filters_down1[-1]] + num_filters_down2
filters_pairs_d2 = [[num_filters[i], num_filters[i + 1]]
for i in range(len(num_filters) - 1)]
for i, o in filters_pairs_d2:
middle_layers.append(
spconv.SubMConv3d(i, o, 3, bias=False, indice_key="subm1"))
middle_layers.append(BatchNorm1d(o))
middle_layers.append(nn.ReLU())
middle_layers.append(
spconv.SparseConv3d(num_filters[-1],
num_filters[-1], (3, 1, 1), (2, 1, 1),
bias=False))
middle_layers.append(BatchNorm1d(num_filters[-1]))
middle_layers.append(nn.ReLU())
self.middle_conv = spconv.SparseSequential(*middle_layers)
def __init__(self, C_in, grid_shape, cfg):
super(CNN_3D, self).__init__()
self.blocks = spconv.SparseSequential(
spconv.SparseConv3d(C_in, 16, 3, 1, padding=0, bias=False),
spconv.SparseConv3d(16, 16, 3, 2, padding=1, bias=False),
spconv.SparseConv3d(16, 32, 3, 2, padding=1, bias=False),
spconv.SparseConv3d(32, 64, 3, 2, padding=1, bias=False),
)
self.grid_shape = grid_shape
self.base_voxel_size = torch.cuda.FloatTensor(cfg.voxel_size)
self.voxel_offset = torch.cuda.FloatTensor(cfg.grid_bounds[:3])
def __init__(self, shape):
super().__init__()
self.net = spconv.SparseConv3d(32, 64, 3)
print(type(self.net.weight))
# self.net.bias = 0.1
# # print(self.net.weight.shape)
# print(self.net.bias)
# print(self.net.bias.)
# exit()
self.sub_net = spconv.SparseConv3d(32, 64, 3)
print(self.sub_net.weight)
exit()
self.shape = shape
def post_act_block(self, in_channels, out_channels, kernel_size, indice_key, stride=1, padding=0,
conv_type='subm', norm_fn=None):
if conv_type == 'subm':
m = spconv.SparseSequential(
spconv.SubMConv3d(in_channels, out_channels, kernel_size, bias=False, indice_key=indice_key),
norm_fn(out_channels),
nn.ReLU(),
)
elif conv_type == 'spconv':
m = spconv.SparseSequential(
spconv.SparseConv3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding,
bias=False, indice_key=indice_key),
norm_fn(out_channels),
nn.ReLU(),
)
elif conv_type == 'inverseconv':
m = spconv.SparseSequential(
spconv.SparseInverseConv3d(in_channels, out_channels, kernel_size,
indice_key=indice_key, bias=False),
norm_fn(out_channels),
nn.ReLU(),
)
else:
raise NotImplementedError
return m
def make_sparse_conv_layer(C_in, C_out, *args, **kwargs):
layer = spconv.SparseSequential(
spconv.SparseConv3d(C_in, C_out, *args, **kwargs),
build_batchnorm(C_out),
nn.ReLU(),
)
return layer
def make_sparse_conv_layer(C_in, C_out, *args, **kwargs):
layer = spconv.SparseSequential(
spconv.SparseConv3d(C_in, C_out, *args, **kwargs, bias=False),
nn.BatchNorm1d(C_out, eps=1e-3, momentum=0.01),
nn.ReLU(),
)
return layer
def __init__(self, nPlanes, norm_fn, block_reps, block, indice_key_id=1):
super().__init__()
self.nPlanes = nPlanes
blocks = {'block{}'.format(i): block(nPlanes[0], nPlanes[0], norm_fn, indice_key='subm{}'.format(indice_key_id)) for i in range(block_reps)}
blocks = OrderedDict(blocks)
self.blocks = spconv.SparseSequential(blocks)
if len(nPlanes) > 1:
self.conv = spconv.SparseSequential(
norm_fn(nPlanes[0]),
nn.ReLU(),
spconv.SparseConv3d(nPlanes[0], nPlanes[1], kernel_size=2, stride=2, bias=False, indice_key='spconv{}'.format(indice_key_id))
)
self.u = UBlock(nPlanes[1:], norm_fn, block_reps, block, indice_key_id=indice_key_id+1)
self.deconv = spconv.SparseSequential(
norm_fn(nPlanes[1]),
nn.ReLU(),
spconv.SparseInverseConv3d(nPlanes[1], nPlanes[0], kernel_size=2, bias=False, indice_key='spconv{}'.format(indice_key_id))
)
blocks_tail = {}
for i in range(block_reps):
blocks_tail['block{}'.format(i)] = block(nPlanes[0] * (2 - i), nPlanes[0], norm_fn, indice_key='subm{}'.format(indice_key_id))
blocks_tail = OrderedDict(blocks_tail)
self.blocks_tail = spconv.SparseSequential(blocks_tail)
def __init__(self, shape):
super().__init__()
self.net = spconv.SparseSequential(
spconv.SparseConv3d(1, 32, 4),
nn.LeakyReLU(),
spconv.SparseConv3d(32, 32, 4),
nn.LeakyReLU(),
spconvpool.SparseMaxPool(3, 5),
spconv.SparseConv3d(32, 32, 4),
nn.LeakyReLU(),
spconv.SparseConv3d(32, 32, 4),
nn.LeakyReLU(),
spconvpool.SparseMaxPool(3, 5),
spconv.SparseConv3d(32, 32, 4),
nn.LeakyReLU(),
spconv.SparseConv3d(32, 32, 4),
nn.LeakyReLU(),
spconv.SparseConv3d(32, 1, 4),
nn.LeakyReLU(),
spconvpool.SparseMaxPool(3, 5),
spconv.ToDense(),
nn.Flatten(),
nn.Linear(14688, 1000),
nn.LeakyReLU(),
nn.Linear(1000, 1000),
nn.LeakyReLU(),
nn.Linear(1000, 1),
nn.LeakyReLU(),
# nn.Sigmoid()
)
self.shape = shape
def stride_conv(in_channels, out_channels, indice_key=None):
return spconv.SparseSequential(
spconv.SparseConv3d(in_channels,
out_channels,
3, (2, 2, 2),
padding=1,
bias=False,
indice_key=indice_key),
nn.BatchNorm1d(out_channels, eps=1e-3, momentum=0.01), nn.ReLU())
def __init__(self, input_channels, grid_size):
super().__init__()
self.sparse_shape = grid_size[::-1] + [1, 0, 0]
self.conv_input = spconv.SparseSequential(
spconv.SparseConv3d(input_channels, 16, 3, padding=1, bias=False, indice_key='subm1'),
nn.BatchNorm1d(16, eps=1e-3, momentum=0.01),
nn.ReLU(),
)
self.conv1 = createSparseConvBlock(16, 16, 3, indice_key='subm1')
self.conv2 = spconv.SparseSequential(
# [1600, 1408, 41] -> [800, 704, 21]
createSparseConvBlock(16, 32, 3, stride=2, indice_key='spconv2', conv_type='spconv'),
createSparseConvBlock(32, 32, 3, indice_key='subm2'),
createSparseConvBlock(32, 32, 3, indice_key='subm2'),
)
self.conv3 = spconv.SparseSequential(
# [800, 704, 21] -> [400, 352, 11]
createSparseConvBlock(32, 64, 3, stride=2, indice_key='spconv3', conv_type='spconv'),
createSparseConvBlock(64, 64, 3, indice_key='subm3'),
createSparseConvBlock(64, 64, 3, indice_key='subm3'),
)
self.conv4 = spconv.SparseSequential(
# [400, 352, 11] -> [200, 176, 5]
createSparseConvBlock(64, 64, 3, stride=2, padding=(0, 1, 1), indice_key='spconv4', conv_type='spconv'),
createSparseConvBlock(64, 64, 3, indice_key='subm4'),
createSparseConvBlock(64, 64, 3, indice_key='subm4'),
)
self.conv_out = spconv.SparseSequential(
# [200, 150, 5] -> [200, 150, 2]
spconv.SparseConv3d(64, 128, (3, 1, 1), stride=(2, 1, 1), padding=0,
bias=False, indice_key='spconv_down2'),
nn.BatchNorm1d(128, eps=1e-3, momentum=0.01),
nn.ReLU(),
)
def __init__(self, model_cfg, input_channels, grid_size, **kwargs):
super().__init__()
self.model_cfg = model_cfg
norm_fn = partial(nn.BatchNorm1d, eps=1e-3, momentum=0.01)
self.sparse_shape = grid_size[::-1] + [1, 0, 0]
self.conv_input = spconv.SparseSequential(
spconv.SubMConv3d(input_channels, 16, 3, padding=1, bias=False, indice_key='subm1'),
norm_fn(16),
nn.ReLU(),
)
block = post_act_block
self.conv1 = spconv.SparseSequential(
SparseBasicBlock(16, 16, norm_fn=norm_fn, indice_key='res1'),
SparseBasicBlock(16, 16, norm_fn=norm_fn, indice_key='res1'),
)
self.conv2 = spconv.SparseSequential(
# [1600, 1408, 41] <- [800, 704, 21]
block(16, 32, 3, norm_fn=norm_fn, stride=2, padding=1, indice_key='spconv2', conv_type='spconv'),
SparseBasicBlock(32, 32, norm_fn=norm_fn, indice_key='res2'),
SparseBasicBlock(32, 32, norm_fn=norm_fn, indice_key='res2'),
)
self.conv3 = spconv.SparseSequential(
# [800, 704, 21] <- [400, 352, 11]
block(32, 64, 3, norm_fn=norm_fn, stride=2, padding=1, indice_key='spconv3', conv_type='spconv'),
SparseBasicBlock(64, 64, norm_fn=norm_fn, indice_key='res3'),
SparseBasicBlock(64, 64, norm_fn=norm_fn, indice_key='res3'),
)
self.conv4 = spconv.SparseSequential(
# [400, 352, 11] <- [200, 176, 5]
block(64, 128, 3, norm_fn=norm_fn, stride=2, padding=(0, 1, 1), indice_key='spconv4', conv_type='spconv'),
SparseBasicBlock(128, 128, norm_fn=norm_fn, indice_key='res4'),
SparseBasicBlock(128, 128, norm_fn=norm_fn, indice_key='res4'),
)
last_pad = 0
last_pad = self.model_cfg.get('last_pad', last_pad)
self.conv_out = spconv.SparseSequential(
# [200, 150, 5] -> [200, 150, 2]
spconv.SparseConv3d(128, 128, (3, 1, 1), stride=(2, 1, 1), padding=last_pad,
bias=False, indice_key='spconv_down2'),
norm_fn(128),
nn.ReLU(),
)
self.num_point_features = 128
def __init__(self, nPlanes, norm_fn, block_reps, block, indice_key_id=1, backbone=False, UNet_Transformer=None):
super().__init__()
self.nPlanes = nPlanes
self.backbone = backbone
self.UNet_Transformer = UNet_Transformer
blocks = {'block{}'.format(i): block(nPlanes[0], nPlanes[0], norm_fn, indice_key='subm{}'.format(indice_key_id))
for i in range(block_reps)}
blocks = OrderedDict(blocks)
self.blocks = spconv.SparseSequential(blocks)
if len(nPlanes) > 1:
self.conv = spconv.SparseSequential(
norm_fn(nPlanes[0]),
nn.ReLU(),
spconv.SparseConv3d(nPlanes[0], nPlanes[1], kernel_size=2, stride=2, bias=False,
indice_key='spconv{}'.format(indice_key_id))
)
self.u = UBlock(nPlanes[1:], norm_fn, block_reps, block, indice_key_id=indice_key_id + 1,
backbone=backbone, UNet_Transformer=self.UNet_Transformer)
self.deconv = spconv.SparseSequential(
norm_fn(nPlanes[1]),
nn.ReLU(),
spconv.SparseInverseConv3d(nPlanes[1], nPlanes[0], kernel_size=2, bias=False,
indice_key='spconv{}'.format(indice_key_id))
)
blocks_tail = {}
for i in range(block_reps):
blocks_tail['block{}'.format(i)] = block(nPlanes[0] * (2 - i), nPlanes[0], norm_fn,
indice_key='subm{}'.format(indice_key_id))
blocks_tail = OrderedDict(blocks_tail)
self.blocks_tail = spconv.SparseSequential(blocks_tail)
elif self.backbone and self.UNet_Transformer['activate']:
self.transformer_encoder = UNetTransformer(
d_model=self.nPlanes[-1],
nhead=self.UNet_Transformer['multi_heads'],
num_encoder_layers=self.UNet_Transformer['num_encoder_layers'],
dim_feedforward=self.UNet_Transformer['dim_feedforward'],
dropout=self.UNet_Transformer['dropout'],
)
def __init__(self, input_channels):
super().__init__()
norm_fn = partial(nn.BatchNorm1d, eps=1e-3, momentum=0.01)
self.conv_input = spconv.SparseSequential(
spconv.SubMConv3d(input_channels, 16, 3, padding=1, bias=False, indice_key='subm1'),
norm_fn(16),
nn.ReLU(),
)
block = self.post_act_block
self.conv1 = spconv.SparseSequential(
block(16, 16, 3, norm_fn=norm_fn, padding=1, indice_key='subm1'),
)
self.conv2 = spconv.SparseSequential(
# [1600, 1408, 41] <- [800, 704, 21]
block(16, 32, 3, norm_fn=norm_fn, stride=2, padding=1, indice_key='spconv2', conv_type='spconv'),
block(32, 32, 3, norm_fn=norm_fn, padding=1, indice_key='subm2'),
block(32, 32, 3, norm_fn=norm_fn, padding=1, indice_key='subm2'),
)
self.conv3 = spconv.SparseSequential(
# [800, 704, 21] <- [400, 352, 11]
block(32, 64, 3, norm_fn=norm_fn, stride=2, padding=1, indice_key='spconv3', conv_type='spconv'),
block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm3'),
block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm3'),
)
self.conv4 = spconv.SparseSequential(
# [400, 352, 11] <- [200, 176, 5]
block(64, 64, 3, norm_fn=norm_fn, stride=2, padding=(0, 1, 1), indice_key='spconv4', conv_type='spconv'),
block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm4'),
block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm4'),
)
last_pad = 0 if cfg.DATA_CONFIG.VOXEL_GENERATOR.VOXEL_SIZE[-1] in [0.1, 0.2] else (1, 0, 0)
self.conv_out = spconv.SparseSequential(
# [200, 150, 5] -> [200, 150, 2]
spconv.SparseConv3d(64, 128, (3, 1, 1), stride=(2, 1, 1), padding=last_pad,
bias=False, indice_key='spconv_down2'),
norm_fn(128),
nn.ReLU(),
)
def createSparseConvBlock(in_channels, out_channels, kernel_size, indice_key=None, stride=1, padding=1,
conv_type='subm'):
if conv_type == 'subm':
conv = spconv.SubMConv3d(in_channels, out_channels, kernel_size, bias=False, indice_key=indice_key)
elif conv_type == 'spconv':
conv = spconv.SparseConv3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding,
bias=False, indice_key=indice_key)
else:
raise NotImplementedError
m = spconv.SparseSequential(
conv,
nn.BatchNorm1d(out_channels, eps=1e-3, momentum=0.01),
nn.ReLU(),
)
return m
def __init__(self, num_layers, ndim, shape, in_channels, out_channels,
kernel_size, stride):
super().__init__()
self.net = spconv.SparseSequential(
spconv.SparseConv3d(in_channels,
out_channels,
kernel_size,
stride,
indice_key="cp0",
bias=False),
spconv.SparseInverseConv3d(out_channels,
in_channels,
kernel_size,
indice_key="cp0",
bias=False),
)
self.todense = spconv.ToDense()
self.shape = shape
def __init__(self, shape): super().__init__() # For some reason the batch normalization is screwing things up idk why. # I heard it fixes itself when you do model.eval rather than just getting a single output self.net = spconv.SparseSequential( spconv.SparseConv3d(1, 32, 4), nn.LeakyReLU(), spconv.SparseConv3d(32, 32, 4), nn.LeakyReLU(), spconvpool.SparseMaxPool(3, 5), spconv.SparseConv3d(32, 32, 4), nn.LeakyReLU(), spconv.SparseConv3d(32, 32, 4), nn.LeakyReLU(), spconvpool.SparseMaxPool(3, 5), spconv.SparseConv3d(32, 32, 4), nn.LeakyReLU(), spconv.SparseConv3d(32, 32, 4), nn.LeakyReLU(), spconv.SparseConv3d(32, 1, 4), nn.LeakyReLU(), spconvpool.SparseMaxPool(3, 5), spconv.ToDense(), nn.Flatten(), nn.Linear(14688, 1000), nn.LeakyReLU(), nn.Linear(1000, 1000), nn.LeakyReLU(), nn.Linear(1000, 1), nn.LeakyReLU(), # nn.Sigmoid() ) self.shape = shape
def generate_block(self,
in_dim,
out_dim,
ksize,
stride,
padding,
do_subm=True):
block = spconv.SparseSequential(
spconv.SubMConv3d(in_channels=in_dim,
out_channels=out_dim,
kernel_size=1,
stride=1,
indice_key="subm0"),
nn.BatchNorm1d(num_features=out_dim),
nn.LeakyReLU(),
spconv.SubMConv3d(in_channels=in_dim,
out_channels=out_dim,
kernel_size=3,
stride=1,
padding=1,
indice_key='subm0')
if do_subm else spconv.SparseConv3d(in_channels=in_dim,
out_channels=out_dim,
kernel_size=3,
stride=1,
padding=1),
nn.BatchNorm1d(num_features=out_dim),
nn.LeakyReLU(),
spconv.SubMConv3d(in_channels=in_dim,
out_channels=out_dim,
kernel_size=1,
stride=1,
indice_key="subm0"),
nn.BatchNorm1d(num_features=out_dim),
spconv.ToDense(),
)
return block
def __init__(self, model_cfg, input_channels, grid_size, voxel_size,
point_cloud_range, **kwargs):
super().__init__()
self.model_cfg = model_cfg
self.sparse_shape = grid_size[::-1] + [1, 0, 0]
self.voxel_size = voxel_size
self.point_cloud_range = point_cloud_range
norm_fn = partial(nn.BatchNorm1d, eps=1e-3, momentum=0.01)
self.conv_input = spconv.SparseSequential(
spconv.SubMConv3d(input_channels,
16,
3,
padding=1,
bias=False,
indice_key='subm1'),
norm_fn(16),
nn.ReLU(),
)
block = post_act_block
self.conv1 = spconv.SparseSequential(
block(16, 16, 3, norm_fn=norm_fn, padding=1, indice_key='subm1'), )
self.conv2 = spconv.SparseSequential(
# [1600, 1408, 41] <- [800, 704, 21]
block(16,
32,
3,
norm_fn=norm_fn,
stride=2,
padding=1,
indice_key='spconv2',
conv_type='spconv'),
block(32, 32, 3, norm_fn=norm_fn, padding=1, indice_key='subm2'),
block(32, 32, 3, norm_fn=norm_fn, padding=1, indice_key='subm2'),
)
self.conv3 = spconv.SparseSequential(
# [800, 704, 21] <- [400, 352, 11]
block(32,
64,
3,
norm_fn=norm_fn,
stride=2,
padding=1,
indice_key='spconv3',
conv_type='spconv'),
block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm3'),
block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm3'),
)
self.conv4 = spconv.SparseSequential(
# [400, 352, 11] <- [200, 176, 5]
block(64,
64,
3,
norm_fn=norm_fn,
stride=2,
padding=(0, 1, 1),
indice_key='spconv4',
conv_type='spconv'),
block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm4'),
block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm4'),
)
if self.model_cfg.get('RETURN_ENCODED_TENSOR', True):
last_pad = self.model_cfg.get('last_pad', 0)
self.conv_out = spconv.SparseSequential(
# [200, 150, 5] -> [200, 150, 2]
spconv.SparseConv3d(64,
128, (3, 1, 1),
stride=(2, 1, 1),
padding=last_pad,
bias=False,
indice_key='spconv_down2'),
norm_fn(128),
nn.ReLU(),
)
else:
self.conv_out = None
# decoder
# [400, 352, 11] <- [200, 176, 5]
self.conv_up_t4 = SparseBasicBlock(64,
64,
indice_key='subm4',
norm_fn=norm_fn)
self.conv_up_m4 = block(128,
64,
3,
norm_fn=norm_fn,
padding=1,
indice_key='subm4')
self.inv_conv4 = block(64,
64,
3,
norm_fn=norm_fn,
indice_key='spconv4',
conv_type='inverseconv')
# [800, 704, 21] <- [400, 352, 11]
self.conv_up_t3 = SparseBasicBlock(64,
64,
indice_key='subm3',
norm_fn=norm_fn)
self.conv_up_m3 = block(128,
64,
3,
norm_fn=norm_fn,
padding=1,
indice_key='subm3')
self.inv_conv3 = block(64,
32,
3,
norm_fn=norm_fn,
indice_key='spconv3',
conv_type='inverseconv')
# [1600, 1408, 41] <- [800, 704, 21]
self.conv_up_t2 = SparseBasicBlock(32,
32,
indice_key='subm2',
norm_fn=norm_fn)
self.conv_up_m2 = block(64, 32, 3, norm_fn=norm_fn, indice_key='subm2')
self.inv_conv2 = block(32,
16,
3,
norm_fn=norm_fn,
indice_key='spconv2',
conv_type='inverseconv')
# [1600, 1408, 41] <- [1600, 1408, 41]
self.conv_up_t1 = SparseBasicBlock(16,
16,
indice_key='subm1',
norm_fn=norm_fn)
self.conv_up_m1 = block(32, 16, 3, norm_fn=norm_fn, indice_key='subm1')
self.conv5 = spconv.SparseSequential(
block(16, 16, 3, norm_fn=norm_fn, padding=1, indice_key='subm1'))
self.num_point_features = 16
def __init__(self, in_channel, config):
super().__init__()
self.print_time = False
norm_fn = partial(nn.BatchNorm1d, eps=1e-3, momentum=0.01)
self.conv_input = spconv.SparseSequential(
spconv.SubMConv3d(in_channel,
16,
3,
bias=False,
padding=1,
indice_key="subm1"),
norm_fn(16),
nn.ReLU(),
)
self.config = config
block = self.convblock
if self.config["add_layers_for_conv3d"]:
self.conv1 = spconv.SparseSequential(
block(16,
16,
3,
padding=1,
norm_fn=norm_fn,
indice_key="subm1"),
block(16,
16,
3,
padding=1,
norm_fn=norm_fn,
indice_key="subm1"),
)
self.conv2 = spconv.SparseSequential(
block(16,
32,
3,
stride=2,
padding=1,
norm_fn=norm_fn,
conv_type="spconv",
indice_key="spconv2"),
block(32,
32,
3,
stride=1,
padding=1,
norm_fn=norm_fn,
conv_type="subm",
indice_key="subm2"),
block(32,
32,
3,
stride=1,
padding=1,
norm_fn=norm_fn,
conv_type="subm",
indice_key="subm2"),
block(32,
32,
3,
stride=1,
padding=1,
norm_fn=norm_fn,
conv_type="subm",
indice_key="subm2"),
block(32,
32,
3,
stride=1,
padding=1,
norm_fn=norm_fn,
conv_type="subm",
indice_key="subm2"),
)
self.conv3 = spconv.SparseSequential(
block(32,
64,
3,
stride=2,
padding=1,
norm_fn=norm_fn,
conv_type="spconv",
indice_key="spconv3"),
block(64,
64,
3,
stride=1,
padding=1,
norm_fn=norm_fn,
conv_type="subm",
indice_key="subm3"),
block(64,
64,
3,
stride=1,
padding=1,
norm_fn=norm_fn,
conv_type="subm",
indice_key="subm3"),
block(64,
64,
3,
stride=1,
padding=1,
norm_fn=norm_fn,
conv_type="subm",
indice_key="subm3"),
block(64,
64,
3,
stride=1,
padding=1,
norm_fn=norm_fn,
conv_type="subm",
indice_key="subm3"),
)
self.conv4 = spconv.SparseSequential(
block(64,
128,
3,
stride=2,
padding=1,
norm_fn=norm_fn,
conv_type="spconv",
indice_key="spconv4"),
block(128,
128,
3,
stride=1,
padding=1,
norm_fn=norm_fn,
conv_type="subm",
indice_key="subm4"),
block(128,
128,
3,
stride=1,
padding=1,
norm_fn=norm_fn,
conv_type="subm",
indice_key="subm4"),
block(128,
128,
3,
stride=1,
padding=1,
norm_fn=norm_fn,
conv_type="subm",
indice_key="subm4"),
block(128,
128,
3,
stride=1,
padding=1,
norm_fn=norm_fn,
conv_type="subm",
indice_key="subm4"),
)
self.conv_out = spconv.SparseSequential(
block(128,
256, (3, 1, 1),
stride=(2, 1, 1),
padding=0,
norm_fn=norm_fn,
conv_type="spconv",
indice_key="conv_down"),
block(256,
256,
3,
stride=1,
padding=1,
norm_fn=norm_fn,
conv_type="subm",
indice_key="subm5"),
block(256,
256,
3,
stride=1,
padding=1,
norm_fn=norm_fn,
conv_type="subm",
indice_key="subm5"),
block(256,
256,
3,
stride=1,
padding=1,
norm_fn=norm_fn,
conv_type="subm",
indice_key="subm5"),
)
else:
self.conv1 = spconv.SparseSequential(
block(16,
16,
3,
padding=1,
norm_fn=norm_fn,
indice_key="subm1"))
self.conv2 = spconv.SparseSequential(
block(16,
32,
3,
stride=2,
padding=1,
norm_fn=norm_fn,
conv_type="spconv",
indice_key="spconv2"),
block(32,
32,
3,
stride=1,
padding=1,
norm_fn=norm_fn,
conv_type="subm",
indice_key="subm2"),
block(32,
32,
3,
stride=1,
padding=1,
norm_fn=norm_fn,
conv_type="subm",
indice_key="subm2"),
block(32,
32,
3,
stride=1,
padding=1,
norm_fn=norm_fn,
conv_type="subm",
indice_key="subm2"),
)
self.conv3 = spconv.SparseSequential(
block(32,
64,
3,
stride=2,
padding=1,
norm_fn=norm_fn,
conv_type="spconv",
indice_key="spconv3"),
block(64,
64,
3,
stride=1,
padding=1,
norm_fn=norm_fn,
conv_type="subm",
indice_key="subm3"),
block(64,
64,
3,
stride=1,
padding=1,
norm_fn=norm_fn,
conv_type="subm",
indice_key="subm3"),
block(64,
64,
3,
stride=1,
padding=1,
norm_fn=norm_fn,
conv_type="subm",
indice_key="subm3"),
)
self.conv4 = spconv.SparseSequential(
block(64,
128,
3,
stride=2,
padding=1,
norm_fn=norm_fn,
conv_type="spconv",
indice_key="spconv4"),
block(128,
128,
3,
stride=1,
padding=1,
norm_fn=norm_fn,
conv_type="subm",
indice_key="subm4"),
block(128,
128,
3,
stride=1,
padding=1,
norm_fn=norm_fn,
conv_type="subm",
indice_key="subm4"),
block(128,
128,
3,
stride=1,
padding=1,
norm_fn=norm_fn,
conv_type="subm",
indice_key="subm4"),
)
self.conv_out = spconv.SparseSequential(
spconv.SparseConv3d(128,
256,
kernel_size=(3, 1, 1),
stride=(2, 1, 1),
padding=0,
indice_key="conv_down"), norm_fn(256),
nn.ReLU())
self.ret = {}
import torch import spconv #from lib.pointgroup_ops.functions import pointgroup_ops torch.jit.script(spconv.SparseConv3d(32, 64, 3)) #torch.jit.script(pointgroup_ops.BFSCluster())
def __init__(self, in_channel, pred_dim, chans=64):
super(SparseNet3D, self).__init__()
conv3d = []
up_bn = [] #batch norm layer for deconvolution
conv3d_transpose = []
# self.conv3d = torch.nn.Conv3d(4, 32, (4,4,4), stride=(2,2,2), padding=(1,1,1))
# self.layers = []
self.down_in_dims = [in_channel, chans, 2 * chans]
self.down_out_dims = [chans, 2 * chans, 4 * chans]
self.down_ksizes = [4, 4, 4]
self.down_strides = [2, 2, 2]
padding = 1 #Note: this only holds for ksize=4 and stride=2!
print('down dims: ', self.down_out_dims)
for i, (in_dim, out_dim, ksize, stride) in enumerate(
zip(self.down_in_dims, self.down_out_dims, self.down_ksizes,
self.down_strides)):
# print('3D CONV', end=' ')
conv3d += [
spconv.SparseConv3d(in_channels=in_dim,
out_channels=out_dim,
kernel_size=ksize,
stride=stride,
padding=padding),
nn.LeakyReLU(),
nn.BatchNorm1d(num_features=out_dim),
]
# self.conv3d = nn.ModuleList(conv3d)
self.conv3d = spconv.SparseSequential(*conv3d)
self.up_in_dims = [4 * chans, 6 * chans]
self.up_bn_dims = [6 * chans, 3 * chans]
self.up_out_dims = [4 * chans, 2 * chans]
self.up_ksizes = [4, 4]
self.up_strides = [2, 2]
padding = 1 #Note: this only holds for ksize=4 and stride=2!
print('up dims: ', self.up_out_dims)
for i, (in_dim, bn_dim, out_dim, ksize, stride) in enumerate(
zip(self.up_in_dims, self.up_bn_dims, self.up_out_dims,
self.up_ksizes, self.up_strides)):
conv3d_transpose += [
spconv.SparseConvTranspose3d(in_channels=in_dim,
out_channels=out_dim,
kernel_size=ksize,
stride=stride,
padding=padding),
nn.LeakyReLU(),
nn.BatchNorm1d(num_features=out_dim),
]
up_bn.append(nn.BatchNorm1d(num_features=bn_dim))
# final 1x1x1 conv to get our desired pred_dim
self.final_feature = spconv.SparseSequential(
spconv.ToDense(),
nn.Conv3d(in_channels=2 * chans,
out_channels=pred_dim,
kernel_size=1,
stride=1,
padding=0))
self.conv3d_transpose = spconv.SparseSequential(*conv3d_transpose)
self.up_bn = nn.ModuleList(up_bn)
def __init__(self, model_cfg, **kwargs):
super(PV_ENcoNet_POOLING_SC3, self).__init__()
self.device = torch.device(
'cuda' if torch.cuda.is_available() else 'cpu')
device = self.device
self.model_cfg = model_cfg
self.num_neighbors = self.model_cfg.NUM_NEIGHBORS
self.decimation = self.model_cfg.DECIMATION
self.d_in = self.model_cfg.D_IN
norm_fn = partial(nn.BatchNorm1d, eps=1e-3, momentum=0.01)
self.use_rgb = self.model_cfg.USE_RGB
self.use_nbg = self.model_cfg.USE_NBG
self.randla_encoder = LocalFeatureAggregation(self.d_in, 16,
self.num_neighbors,
device, self.use_rgb)
## (B,N,32,1)
self.sparse_shape = [41, 1600, 1408]
self.conv_input = spconv.SparseSequential(
spconv.SubMConv3d(38,
32,
3,
padding=1,
bias=False,
indice_key='subm1'),
norm_fn(32),
nn.ReLU(),
)
block = post_act_block
self.conv1 = spconv.SparseSequential(
block(32, 32, 3, norm_fn=norm_fn, padding=1, indice_key='subm1'), )
self.conv2 = spconv.SparseSequential(
# [1600, 1408, 41] <- [800, 704, 21]
block(32,
48,
3,
norm_fn=norm_fn,
stride=2,
padding=1,
indice_key='spconv2',
conv_type='spconv'),
block(48, 48, 3, norm_fn=norm_fn, padding=1, indice_key='subm2'),
block(48, 48, 3, norm_fn=norm_fn, padding=1, indice_key='subm2'),
)
self.conv3 = spconv.SparseSequential(
# [800, 704, 21] <- [400, 352, 11]
block(48,
64,
3,
norm_fn=norm_fn,
stride=2,
padding=1,
indice_key='spconv3',
conv_type='spconv'),
block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm3'),
block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm3'),
)
self.conv4 = spconv.SparseSequential(
# [400, 352, 11] <- [200, 176, 5]
block(64,
64,
3,
norm_fn=norm_fn,
stride=2,
padding=(0, 1, 1),
indice_key='spconv4',
conv_type='spconv'),
block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm4'),
block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm4'),
)
last_pad = 0
last_pad = self.model_cfg.get('last_pad', last_pad)
self.conv_out = spconv.SparseSequential(
# [200, 150, 5] -> [200, 150, 2]
spconv.SparseConv3d(64,
128, (3, 1, 1),
stride=(2, 1, 1),
padding=last_pad,
bias=False,
indice_key='spconv_down2'),
norm_fn(128),
nn.ReLU(),
)
self.device = device
self.num_point_features = 128
self = self.to(device)
def __init__(self,
sparse_shape,
heads,
head_conv,
num_input_features=4,
bias=False):
super(CenterNet3D, self).__init__()
self.sparse_shape = sparse_shape
self.heads = heads
self.net3d = spconv.SparseSequential(
spconv.SubMConv3d(num_input_features, 16, 3, 1, 1, bias=bias),
nn.BatchNorm1d(16), nn.LeakyReLU(inplace=True),
spconv.SubMConv3d(16, 16, 3, 1, 1, bias=bias), nn.BatchNorm1d(16),
nn.LeakyReLU(inplace=True),
spconv.SparseConv3d(16, 32, 3, 2, 1,
bias=bias), nn.BatchNorm1d(32),
nn.LeakyReLU(inplace=True),
spconv.SubMConv3d(32, 32, 3, 1, 1, bias=bias), nn.BatchNorm1d(32),
nn.LeakyReLU(inplace=True),
spconv.SubMConv3d(32, 32, 3, 1, 1, bias=bias), nn.BatchNorm1d(32),
nn.LeakyReLU(inplace=True),
spconv.SparseConv3d(32, 64, 3, 2, 1,
bias=bias), nn.BatchNorm1d(64),
nn.LeakyReLU(inplace=True),
spconv.SubMConv3d(64, 64, 3, 1, 1, bias=bias), nn.BatchNorm1d(64),
nn.LeakyReLU(inplace=True),
spconv.SubMConv3d(64, 64, 3, 1, 1, bias=bias), nn.BatchNorm1d(64),
nn.LeakyReLU(inplace=True),
spconv.SubMConv3d(64, 64, 3, 1, 1, bias=bias), nn.BatchNorm1d(64),
nn.LeakyReLU(inplace=True),
spconv.SparseConv3d(64, 64, 3, 2, 1, bias=bias),
nn.BatchNorm1d(64), nn.LeakyReLU(inplace=True),
spconv.SubMConv3d(64, 64, 3, 1, 1, bias=bias), nn.BatchNorm1d(64),
nn.LeakyReLU(inplace=True),
spconv.SubMConv3d(64, 64, 3, 1, 1, bias=bias), nn.BatchNorm1d(64),
nn.LeakyReLU(inplace=True),
spconv.SubMConv3d(64, 64, 3, 1, 1, bias=bias), nn.BatchNorm1d(64),
nn.LeakyReLU(inplace=True))
self.net2d = nn.Sequential(
nn.Conv2d(320, 128, 3, 1, 1), nn.BatchNorm2d(128),
nn.LeakyReLU(inplace=True), nn.Conv2d(128, 128, 3, 1, 1),
nn.BatchNorm2d(128), nn.LeakyReLU(inplace=True),
nn.Conv2d(128, 128, 3, 1, 1), nn.BatchNorm2d(128),
nn.LeakyReLU(inplace=True), DeformConv2d(128, 128, 3, 1, 1),
nn.BatchNorm2d(128), nn.LeakyReLU(inplace=True),
nn.ConvTranspose2d(128, 128, 3, 2, padding=1, output_padding=1),
nn.BatchNorm2d(128), nn.LeakyReLU(inplace=True))
for head in sorted(self.heads):
num_output = self.heads[head]
fc = nn.Sequential(
nn.Conv2d(128, head_conv, kernel_size=3, padding=1, bias=True),
nn.LeakyReLU(inplace=True),
nn.Conv2d(head_conv,
num_output,
kernel_size=1,
stride=1,
padding=0))
self.__setattr__('head_{}'.format(head), fc)
def U(nPlanes): #Recursive function
m = spconv.SparseSequential()
if len(nPlanes) == 1:
for _ in range(reps):
m.add(
spconv.SparseBasicBlock(nPlanes[0],
nPlanes[0],
3,
indice_key="subm{}".format(
len(nPlanes))))
else:
m = spconv.SparseSequential()
for _ in range(reps):
m.add(
spconv.SparseBasicBlock(nPlanes[0],
nPlanes[0],
3,
indice_key="subm{}".format(
len(nPlanes))))
m.add(spconv.ConcatTable().add(spconv.Identity()).add(
spconv.SparseSequential().add(
spconv.SparseConv3d(
nPlanes[0],
nPlanes[1],
downsample[0],
stride=downsample[1],
bias=False,
indice_key="conv{}".format(len(nPlanes)))).add(
nn.BatchNorm1d(
nPlanes[1], eps=1e-3, momentum=0.01)).add(
nn.ReLU()).add(U(nPlanes[1:])).add(
spconv.SparseInverseConv3d(
nPlanes[1],
nPlanes[0],
downsample[0],
bias=False,
indice_key="conv{}".format(
len(nPlanes)))).add(
nn.BatchNorm1d(
nPlanes[0],
eps=1e-3,
momentum=0.01)).add(
nn.ReLU())))
m.add(spconv.JoinTable())
for i in range(reps):
m.add(
spconv.SubMConv3d(
nPlanes[0] * 2,
nPlanes[0],
3,
bias=False,
indice_key="end_pp{}".format(len(nPlanes)))).add(
nn.BatchNorm1d(nPlanes[0], eps=1e-3,
momentum=0.01)).add(nn.ReLU())
m.add(
spconv.SparseBasicBlock(nPlanes[0],
nPlanes[0],
3,
indice_key="end_pp{}".format(
len(nPlanes))))
return m
def __init__(self, in_channels):
super().__init__()
# self.print_info = cfg.print_info
norm_fn = partial(nn.BatchNorm1d, eps=1e-3, momentum=0.01)
self.conv_input = spconv.SparseSequential(
spconv.SubMConv3d(in_channels,
16,
3,
padding=1,
bias=False,
indice_key='subm1'),
norm_fn(16),
nn.ReLU(),
)
block = self.post_act_block
self.conv1 = spconv.SparseSequential(
block(16, 16, 3, norm_fn=norm_fn, padding=1, indice_key='subm1'), )
self.conv2 = spconv.SparseSequential(
# [1600, 1408, 41] <- [800, 704, 21]
block(16,
32,
3,
norm_fn=norm_fn,
stride=2,
padding=1,
indice_key='spconv2',
conv_type='spconv'),
block(32, 32, 3, norm_fn=norm_fn, padding=1, indice_key='subm2'),
block(32, 32, 3, norm_fn=norm_fn, padding=1, indice_key='subm2'),
)
self.conv3 = spconv.SparseSequential(
# [800, 704, 21] <- [400, 352, 11]
block(32,
64,
3,
norm_fn=norm_fn,
stride=2,
padding=1,
indice_key='spconv3',
conv_type='spconv'),
block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm3'),
block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm3'),
)
self.conv4 = spconv.SparseSequential(
# [400, 352, 11] <- [200, 176, 5]
block(64,
64,
3,
norm_fn=norm_fn,
stride=2,
padding=(0, 1, 1),
indice_key='spconv4',
conv_type='spconv'),
block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm4'),
block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm4'),
)
last_pad = 0 if cfg.DATA_CONFIG.VOXEL_GENERATOR.VOXEL_SIZE[-1] in [
0.1, 0.2
] else (1, 0, 0)
self.conv_out = spconv.SparseSequential(
# [200, 150, 5] -> [200, 150, 2]
spconv.SparseConv3d(64,
128, (3, 1, 1),
stride=(2, 1, 1),
padding=last_pad,
bias=False,
indice_key='spconv_down2'),
norm_fn(128),
nn.ReLU(),
)
self.seg_net = SegNet(16)
self.seg_loss_func = loss_utils.SigmoidFocalClassificationLoss_v1(
alpha=0.25, gamma=2.0)
self.num_point_features = 128
self.ret = {}
def __init__(self, input_channels, **kwargs):
super().__init__(unet_target_cfg=cfg.MODEL.RPN.BACKBONE.TARGET_CONFIG)
norm_fn = partial(nn.BatchNorm1d, eps=1e-3, momentum=0.01)
self.conv_input = spconv.SparseSequential(
spconv.SubMConv3d(input_channels,
16,
3,
padding=1,
bias=False,
indice_key='subm1'),
norm_fn(16),
nn.ReLU(),
)
block = self.post_act_block
self.conv1 = spconv.SparseSequential(
block(16, 16, 3, norm_fn=norm_fn, padding=1, indice_key='subm1'), )
self.conv2 = spconv.SparseSequential(
# [1600, 1408, 41] <- [800, 704, 21]
block(16,
32,
3,
norm_fn=norm_fn,
stride=2,
padding=1,
indice_key='spconv2',
conv_type='spconv'),
block(32, 32, 3, norm_fn=norm_fn, padding=1, indice_key='subm2'),
block(32, 32, 3, norm_fn=norm_fn, padding=1, indice_key='subm2'),
)
self.conv3 = spconv.SparseSequential(
# [800, 704, 21] <- [400, 352, 11]
block(32,
64,
3,
norm_fn=norm_fn,
stride=2,
padding=1,
indice_key='spconv3',
conv_type='spconv'),
block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm3'),
block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm3'),
)
self.conv4 = spconv.SparseSequential(
# [400, 352, 11] <- [200, 176, 5]
block(64,
64,
3,
norm_fn=norm_fn,
stride=2,
padding=(0, 1, 1),
indice_key='spconv4',
conv_type='spconv'),
block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm4'),
block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm4'),
)
last_pad = 0 if cfg.DATA_CONFIG.VOXEL_GENERATOR.VOXEL_SIZE[-1] in [
0.1, 0.2
] else (1, 0, 0)
self.conv_out = spconv.SparseSequential(
# [200, 150, 5] -> [200, 150, 2]
spconv.SparseConv3d(64,
128, (3, 1, 1),
stride=(2, 1, 1),
padding=last_pad,
bias=False,
indice_key='spconv_down2'),
norm_fn(128),
nn.ReLU(),
)
# decoder
# [400, 352, 11] <- [200, 176, 5]
self.conv_up_t4 = SparseBasicBlock(64,
64,
indice_key='subm4',
norm_fn=norm_fn)
self.conv_up_m4 = block(128,
64,
3,
norm_fn=norm_fn,
padding=1,
indice_key='subm4')
self.inv_conv4 = block(64,
64,
3,
norm_fn=norm_fn,
indice_key='spconv4',
conv_type='inverseconv')
# [800, 704, 21] <- [400, 352, 11]
self.conv_up_t3 = SparseBasicBlock(64,
64,
indice_key='subm3',
norm_fn=norm_fn)
self.conv_up_m3 = block(128,
64,
3,
norm_fn=norm_fn,
padding=1,
indice_key='subm3')
self.inv_conv3 = block(64,
32,
3,
norm_fn=norm_fn,
indice_key='spconv3',
conv_type='inverseconv')
# [1600, 1408, 41] <- [800, 704, 21]
self.conv_up_t2 = SparseBasicBlock(32,
32,
indice_key='subm2',
norm_fn=norm_fn)
self.conv_up_m2 = block(64, 32, 3, norm_fn=norm_fn, indice_key='subm2')
self.inv_conv2 = block(32,
16,
3,
norm_fn=norm_fn,
indice_key='spconv2',
conv_type='inverseconv')
# [1600, 1408, 41] <- [1600, 1408, 41]
self.conv_up_t1 = SparseBasicBlock(16,
16,
indice_key='subm1',
norm_fn=norm_fn)
self.conv_up_m1 = block(32, 16, 3, norm_fn=norm_fn, indice_key='subm1')
self.conv5 = spconv.SparseSequential(
block(16, 16, 3, norm_fn=norm_fn, padding=1, indice_key='subm1'))
self.seg_cls_layer = nn.Linear(16, 1, bias=True)
self.seg_reg_layer = nn.Linear(16, 3, bias=True)
