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 '''