def __init__(self, in_channels, out_channels, norm_fn, indice_key=None):
super().__init__()
if in_channels == out_channels:
self.i_branch = spconv.SparseSequential(nn.Identity())
else:
self.i_branch = spconv.SparseSequential(
spconv.SubMConv3d(in_channels,
out_channels,
kernel_size=1,
bias=False))
self.conv_branch = spconv.SparseSequential(
norm_fn(in_channels), nn.ReLU(),
spconv.SubMConv3d(in_channels,
out_channels,
kernel_size=3,
padding=1,
bias=False,
indice_key=indice_key), norm_fn(out_channels),
nn.ReLU(),
spconv.SubMConv3d(out_channels,
out_channels,
kernel_size=3,
padding=1,
bias=False,
indice_key=indice_key))
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 make_encoder_layers(self, make_block, norm_cfg, in_channels):
"""make encoder layers using sparse convs.
Args:
make_block (method): A bounded function to build blocks.
norm_cfg (dict[str]): Config of normalization layer.
in_channels (int): The number of encoder input channels.
Returns:
int: The number of encoder output channels.
"""
self.encoder_layers = spconv.SparseSequential()
for i, blocks in enumerate(self.encoder_channels):
blocks_list = []
for j, out_channels in enumerate(tuple(blocks)):
padding = tuple(self.encoder_paddings[i])[j]
# each stage started with a spconv layer
# except the first stage
if i != 0 and j == 0:
blocks_list.append(
make_block(
in_channels,
out_channels,
3,
norm_cfg=norm_cfg,
stride=2,
padding=padding,
indice_key=f'spconv{i + 1}',
conv_type='SparseConv3d'))
else:
blocks_list.append(
make_block(
in_channels,
out_channels,
3,
norm_cfg=norm_cfg,
padding=padding,
indice_key=f'subm{i + 1}',
conv_type='SubMConv3d'))
in_channels = out_channels
stage_name = f'encoder_layer{i + 1}'
stage_layers = spconv.SparseSequential(*blocks_list)
self.encoder_layers.add_module(stage_name, stage_layers)
return out_channels
def __init__(self, in_channels, out_channels, norm_fn, indice_key=None):
super().__init__()
self.conv_layers = spconv.SparseSequential(
norm_fn(in_channels), nn.ReLU(),
spconv.SubMConv3d(in_channels,
out_channels,
kernel_size=3,
padding=1,
bias=False,
indice_key=indice_key))
def __init__(self,
num_classes,
seg_in_channels,
part_in_channels,
seg_conv_channels=None,
part_conv_channels=None,
merge_conv_channels=None,
down_conv_channels=None,
shared_fc_channels=None,
cls_channels=None,
reg_channels=None,
dropout_ratio=0.1,
roi_feat_size=14,
with_corner_loss=True,
bbox_coder=dict(type='DeltaXYZWLHRBBoxCoder'),
conv_cfg=dict(type='Conv1d'),
norm_cfg=dict(type='BN1d', eps=1e-3, momentum=0.01),
loss_bbox=dict(type='SmoothL1Loss',
beta=1.0 / 9.0,
loss_weight=2.0),
loss_cls=dict(type='CrossEntropyLoss',
use_sigmoid=True,
reduction='none',
loss_weight=1.0)):
super(PartA2BboxHead, self).__init__()
self.num_classes = num_classes
self.with_corner_loss = with_corner_loss
self.bbox_coder = build_bbox_coder(bbox_coder)
self.loss_bbox = build_loss(loss_bbox)
self.loss_cls = build_loss(loss_cls)
self.use_sigmoid_cls = loss_cls.get('use_sigmoid', False)
assert down_conv_channels[-1] == shared_fc_channels[0]
# init layers
part_channel_last = part_in_channels
part_conv = []
for i, channel in enumerate(part_conv_channels):
part_conv.append(
make_sparse_convmodule(part_channel_last,
channel,
3,
padding=1,
norm_cfg=norm_cfg,
indice_key=f'rcnn_part{i}',
conv_type='SubMConv3d'))
part_channel_last = channel
self.part_conv = spconv.SparseSequential(*part_conv)
seg_channel_last = seg_in_channels
seg_conv = []
for i, channel in enumerate(seg_conv_channels):
seg_conv.append(
make_sparse_convmodule(seg_channel_last,
channel,
3,
padding=1,
norm_cfg=norm_cfg,
indice_key=f'rcnn_seg{i}',
conv_type='SubMConv3d'))
seg_channel_last = channel
self.seg_conv = spconv.SparseSequential(*seg_conv)
self.conv_down = spconv.SparseSequential()
merge_conv_channel_last = part_channel_last + seg_channel_last
merge_conv = []
for i, channel in enumerate(merge_conv_channels):
merge_conv.append(
make_sparse_convmodule(merge_conv_channel_last,
channel,
3,
padding=1,
norm_cfg=norm_cfg,
indice_key='rcnn_down0'))
merge_conv_channel_last = channel
down_conv_channel_last = merge_conv_channel_last
conv_down = []
for i, channel in enumerate(down_conv_channels):
conv_down.append(
make_sparse_convmodule(down_conv_channel_last,
channel,
3,
padding=1,
norm_cfg=norm_cfg,
indice_key='rcnn_down1'))
down_conv_channel_last = channel
self.conv_down.add_module('merge_conv',
spconv.SparseSequential(*merge_conv))
self.conv_down.add_module(
'max_pool3d', spconv.SparseMaxPool3d(kernel_size=2, stride=2))
self.conv_down.add_module('down_conv',
spconv.SparseSequential(*conv_down))
shared_fc_list = []
pool_size = roi_feat_size // 2
pre_channel = shared_fc_channels[0] * pool_size**3
for k in range(1, len(shared_fc_channels)):
shared_fc_list.append(
ConvModule(pre_channel,
shared_fc_channels[k],
1,
padding=0,
conv_cfg=conv_cfg,
norm_cfg=norm_cfg,
inplace=True))
pre_channel = shared_fc_channels[k]
if k != len(shared_fc_channels) - 1 and dropout_ratio > 0:
shared_fc_list.append(nn.Dropout(dropout_ratio))
self.shared_fc = nn.Sequential(*shared_fc_list)
# Classification layer
channel_in = shared_fc_channels[-1]
cls_channel = 1
cls_layers = []
pre_channel = channel_in
for k in range(0, len(cls_channels)):
cls_layers.append(
ConvModule(pre_channel,
cls_channels[k],
1,
padding=0,
conv_cfg=conv_cfg,
norm_cfg=norm_cfg,
inplace=True))
pre_channel = cls_channels[k]
cls_layers.append(
ConvModule(pre_channel,
cls_channel,
1,
padding=0,
conv_cfg=conv_cfg,
act_cfg=None))
if dropout_ratio >= 0:
cls_layers.insert(1, nn.Dropout(dropout_ratio))
self.conv_cls = nn.Sequential(*cls_layers)
# Regression layer
reg_layers = []
pre_channel = channel_in
for k in range(0, len(reg_channels)):
reg_layers.append(
ConvModule(pre_channel,
reg_channels[k],
1,
padding=0,
conv_cfg=conv_cfg,
norm_cfg=norm_cfg,
inplace=True))
pre_channel = reg_channels[k]
reg_layers.append(
ConvModule(pre_channel,
self.bbox_coder.code_size,
1,
padding=0,
conv_cfg=conv_cfg,
act_cfg=None))
if dropout_ratio >= 0:
reg_layers.insert(1, nn.Dropout(dropout_ratio))
self.conv_reg = nn.Sequential(*reg_layers)
self.init_weights()
def make_sparse_convmodule(in_channels,
out_channels,
kernel_size,
indice_key,
stride=1,
padding=0,
conv_type='SubMConv3d',
norm_cfg=None,
order=('conv', 'norm', 'act')):
"""Make sparse convolution module.
Args:
in_channels (int): the number of input channels
out_channels (int): the number of out channels
kernel_size (int|tuple(int)): kernel size of convolution
indice_key (str): the indice key used for sparse tensor
stride (int|tuple(int)): the stride of convolution
padding (int or list[int]): the padding number of input
conv_type (str): sparse conv type in spconv
norm_cfg (dict[str]): config of normalization layer
order (tuple[str]): The order of conv/norm/activation layers. It is a
sequence of "conv", "norm" and "act". Common examples are
("conv", "norm", "act") and ("act", "conv", "norm").
Returns:
spconv.SparseSequential: sparse convolution module.
"""
assert isinstance(order, tuple) and len(order) <= 3
assert set(order) | {'conv', 'norm', 'act'} == {'conv', 'norm', 'act'}
conv_cfg = dict(type=conv_type, indice_key=indice_key)
layers = list()
for layer in order:
if layer == 'conv':
if conv_type not in [
'SparseInverseConv3d', 'SparseInverseConv2d',
'SparseInverseConv1d'
]:
layers.append(
build_conv_layer(
conv_cfg,
in_channels,
out_channels,
kernel_size,
stride=stride,
padding=padding,
bias=False))
else:
layers.append(
build_conv_layer(
conv_cfg,
in_channels,
out_channels,
kernel_size,
bias=False))
elif layer == 'norm':
layers.append(build_norm_layer(norm_cfg, out_channels)[1])
elif layer == 'act':
layers.append(nn.ReLU(inplace=True))
layers = spconv.SparseSequential(*layers)
return layers
def make_encoder_layers(self,
make_block,
norm_cfg,
in_channels,
block_type='conv_module',
conv_cfg=dict(type='SubMConv3d')):
"""make encoder layers using sparse convs.
Args:
make_block (method): A bounded function to build blocks.
norm_cfg (dict[str]): Config of normalization layer.
in_channels (int): The number of encoder input channels.
block_type (str): Type of the block to use. Defaults to
'conv_module'.
conv_cfg (dict): Config of conv layer. Defaults to
dict(type='SubMConv3d').
Returns:
int: The number of encoder output channels.
"""
assert block_type in ['conv_module', 'basicblock']
self.encoder_layers = spconv.SparseSequential()
for i, blocks in enumerate(self.encoder_channels):
blocks_list = []
for j, out_channels in enumerate(tuple(blocks)):
padding = tuple(self.encoder_paddings[i])[j]
# each stage started with a spconv layer
# except the first stage
if i != 0 and j == 0 and block_type == 'conv_module':
blocks_list.append(
make_block(in_channels,
out_channels,
3,
norm_cfg=norm_cfg,
stride=2,
padding=padding,
indice_key=f'spconv{i + 1}',
conv_type='SparseConv3d'))
elif block_type == 'basicblock':
if j == len(blocks) - 1 and i != len(
self.encoder_channels) - 1:
blocks_list.append(
make_block(in_channels,
out_channels,
3,
norm_cfg=norm_cfg,
stride=2,
padding=padding,
indice_key=f'spconv{i + 1}',
conv_type='SparseConv3d'))
else:
blocks_list.append(
SparseBasicBlock(out_channels,
out_channels,
norm_cfg=norm_cfg,
conv_cfg=conv_cfg))
else:
blocks_list.append(
make_block(in_channels,
out_channels,
3,
norm_cfg=norm_cfg,
padding=padding,
indice_key=f'subm{i + 1}',
conv_type='SubMConv3d'))
in_channels = out_channels
stage_name = f'encoder_layer{i + 1}'
stage_layers = spconv.SparseSequential(*blocks_list)
self.encoder_layers.add_module(stage_name, stage_layers)
return out_channels
def __init__(self, cfg):
super().__init__()
input_c = cfg.input_channel
m = cfg.m
classes = cfg.classes
block_reps = cfg.block_reps
block_residual = cfg.block_residual
self.cluster_radius = cfg.cluster_radius
self.cluster_meanActive = cfg.cluster_meanActive
self.cluster_shift_meanActive = cfg.cluster_shift_meanActive
self.cluster_npoint_thre = cfg.cluster_npoint_thre
self.score_scale = cfg.score_scale
self.score_fullscale = cfg.score_fullscale
self.mode = cfg.score_mode
self.prepare_epochs = cfg.prepare_epochs
self.pretrain_path = cfg.pretrain_path
self.pretrain_module = cfg.pretrain_module
self.fix_module = cfg.fix_module
norm_fn = functools.partial(nn.BatchNorm1d, eps=1e-4, momentum=0.1)
if block_residual:
block = ResidualBlock
else:
block = VGGBlock
if cfg.use_coords:
input_c += 3
#### backbone
self.input_conv = spconv.SparseSequential(
spconv.SubMConv3d(input_c, m, kernel_size=3, padding=1, bias=False, indice_key='subm1')
)
self.unet = UBlock([m, 2 * m, 3 * m, 4 * m, 5 * m, 6 * m, 7 * m], norm_fn, block_reps, block, indice_key_id=1)
self.output_layer = spconv.SparseSequential(
norm_fn(m),
nn.ReLU()
)
#### semantic segmentation
self.linear = nn.Linear(m, classes) # bias(default): True
#### offset
self.offset = nn.Sequential(
nn.Linear(m, m, bias=True),
norm_fn(m),
nn.ReLU()
)
self.offset_linear = nn.Linear(m, 3, bias=True)
#### score branch
self.score_unet = UBlock([m, 2 * m], norm_fn, 2, block, indice_key_id=1)
self.score_outputlayer = spconv.SparseSequential(
norm_fn(m),
nn.ReLU()
)
self.score_linear = nn.Linear(m, 1)
self.apply(self.set_bn_init)
#### fix parameter
module_map = {'input_conv': self.input_conv, 'unet': self.unet, 'output_layer': self.output_layer,
'linear': self.linear, 'offset': self.offset, 'offset_linear': self.offset_linear,
'score_unet': self.score_unet, 'score_outputlayer': self.score_outputlayer,
'score_linear': self.score_linear}
for m in self.fix_module:
mod = module_map[m]
for param in mod.parameters():
param.requires_grad = False
'''