def __init__(self, model_cfg, voxel_size, point_cloud_range, num_bev_features=None,
num_rawpoint_features=None, **kwargs):
super().__init__()
self.model_cfg = model_cfg
self.voxel_size = voxel_size
self.point_cloud_range = point_cloud_range
SA_cfg = self.model_cfg.SA_LAYER
self.SA_layers = nn.ModuleList()
self.SA_layer_names = []
self.downsample_times_map = {}
c_in = 0
for src_name in self.model_cfg.FEATURES_SOURCE:
if src_name in ['bev', 'raw_points']:
# print(src_name)
continue
self.downsample_times_map[src_name] = SA_cfg[src_name].DOWNSAMPLE_FACTOR
mlps = SA_cfg[src_name].MLPS
for k in range(len(mlps)):
mlps[k] = [mlps[k][0]] + mlps[k]
cur_layer = pointnet2_stack_modules.StackSAModuleMSG(
radii=SA_cfg[src_name].POOL_RADIUS,
nsamples=SA_cfg[src_name].NSAMPLE,
mlps=mlps,
use_xyz=True,
pool_method='max_pool',
)
self.SA_layers.append(cur_layer)
self.SA_layer_names.append(src_name)
c_in += sum([x[-1] for x in mlps])
if 'bev' in self.model_cfg.FEATURES_SOURCE:
c_bev = num_bev_features
c_in += c_bev
if 'raw_points' in self.model_cfg.FEATURES_SOURCE:
mlps = SA_cfg['raw_points'].MLPS
for k in range(len(mlps)):
mlps[k] = [num_rawpoint_features - 3] + mlps[k]
self.SA_rawpoints = pointnet2_stack_modules.StackSAModuleMSG(
radii=SA_cfg['raw_points'].POOL_RADIUS,
nsamples=SA_cfg['raw_points'].NSAMPLE,
mlps=mlps,
use_xyz=True,
pool_method='max_pool'
)
c_in += sum([x[-1] for x in mlps])
self.vsa_point_feature_fusion = nn.Sequential(
nn.Linear(c_in, self.model_cfg.NUM_OUTPUT_FEATURES, bias=False),
nn.BatchNorm1d(self.model_cfg.NUM_OUTPUT_FEATURES),
nn.ReLU(),
)
self.num_point_features = self.model_cfg.NUM_OUTPUT_FEATURES
self.num_point_features_before_fusion = c_in
self.print_info = cfg.print_info
def __init__(self, num_voxel_center_features=3, num_bev_features=256):
super().__init__()
self.features_src = cfg.MODEL.SA_module["feature_source"]
self.conv3d_names = cfg.MODEL.SA_module["conv3d_names"]
self.SA_layer = nn.ModuleList()
self.conv3d_strides = {}
self.point_cloud_range = np.array(cfg.DATA_CONFIG.POINT_CLOUD_RANGE)
self.voxel_size = cfg.DATA_CONFIG.VOXEL_GENERATOR["VOXEL_SIZE"]
self.using_voxel_center = cfg.MODEL.SA_module["using_voxel_center"]
channel_in = 0
if "voxel_centers" in self.features_src:
mlps = cfg.MODEL.SA_module.sa_layer["voxel_centers"]["mlps"]
for i in range(len(mlps)):
if self.using_voxel_center:
mlps[i] = [num_voxel_center_features - 3] + mlps[i]
else:
mlps[i] = [num_voxel_center_features - 3 + 1] + mlps[i]
nsamples = cfg.MODEL.SA_module.sa_layer.voxel_centers["nsamples"]
pool_radius = cfg.MODEL.SA_module.sa_layer.voxel_centers[
"pool_radius"]
self.SA_voxel_centers = pointnet2.StackSAModuleMSG(
nsamples=nsamples,
radii=pool_radius,
mlps=mlps,
use_xyz=True,
pool_method="max_pool")
channel_in += sum([x[-1] for x in mlps])
if "conv_3d" in self.features_src:
for feature_name in self.conv3d_names:
if (feature_name in ["voxel_centers", "bev"]):
continue
mlps = cfg.MODEL.SA_module.sa_layer[feature_name]["mlps"]
self.conv3d_strides[
feature_name] = cfg.MODEL.SA_module.sa_layer[feature_name][
"down_sample_fraction"]
for k in range(len(mlps)):
mlps[k] = [mlps[k][0]] + mlps[k]
nsamples = cfg.MODEL.SA_module.sa_layer[feature_name][
"nsamples"]
pool_radius = cfg.MODEL.SA_module.sa_layer[feature_name][
"pool_radius"]
cur_sa_layer = pointnet2.StackSAModuleMSG(
nsamples=nsamples,
radii=pool_radius,
mlps=mlps,
use_xyz=True,
pool_method="max_pool")
self.SA_layer.append(cur_sa_layer)
channel_in += sum([x[-1] for x in mlps])
if "bev" in self.features_src:
channel_in += num_bev_features
feature_src_out_channel = cfg.MODEL.SA_module[
"point_features_out_channel"]
self.features_src_fusion = nn.Sequential(
nn.Linear(channel_in, feature_src_out_channel, bias=False),
nn.BatchNorm1d(feature_src_out_channel), nn.ReLU())
self.points_out_channels = feature_src_out_channel
self.points_out_channels_befor_fusion = channel_in
self.ret = {}
def __init__(self, input_channels, model_cfg, num_class=1):
super().__init__(num_class=num_class, model_cfg=model_cfg)
self.model_cfg = model_cfg
mlps = self.model_cfg.ROI_GRID_POOL.MLPS
for k in range(len(mlps)):
mlps[k] = [input_channels] + mlps[k]
self.roi_grid_pool_layer = pointnet2_stack_modules.StackSAModuleMSG(
radii=self.model_cfg.ROI_GRID_POOL.POOL_RADIUS,
nsamples=self.model_cfg.ROI_GRID_POOL.NSAMPLE,
mlps=mlps,
use_xyz=True,
pool_method=self.model_cfg.ROI_GRID_POOL.POOL_METHOD,
)
GRID_SIZE = self.model_cfg.ROI_GRID_POOL.GRID_SIZE
c_out = sum([x[-1] for x in mlps])
pre_channel = GRID_SIZE * GRID_SIZE * GRID_SIZE * c_out
shared_fc_list = []
for k in range(0, self.model_cfg.SHARED_FC.__len__()):
shared_fc_list.extend([
nn.Conv1d(pre_channel, self.model_cfg.SHARED_FC[k], kernel_size=1, bias=False),
nn.BatchNorm1d(self.model_cfg.SHARED_FC[k]),
nn.ReLU()
])
pre_channel = self.model_cfg.SHARED_FC[k]
if k != self.model_cfg.SHARED_FC.__len__() - 1 and self.model_cfg.DP_RATIO > 0:
shared_fc_list.append(nn.Dropout(self.model_cfg.DP_RATIO))
self.shared_fc_layer = nn.Sequential(*shared_fc_list)
self.cls_layers = self.make_fc_layers(
input_channels=pre_channel, output_channels=self.num_class, fc_list=self.model_cfg.CLS_FC
)
self.reg_layers = self.make_fc_layers(
input_channels=pre_channel,
output_channels=self.box_coder.code_size * self.num_class,
fc_list=self.model_cfg.REG_FC
)
self.init_weights(weight_init='xavier')
def __init__(self, input_channels):
super().__init__()
self.register_buffer('global_step', torch.LongTensor(1).zero_())
#######################################稀疏三维卷积#########################################
norm_fn = partial(nn.BatchNorm1d, eps=1e-3, momentum=0.01)
self.downsample_times_map = [1, 2, 4, 8]
self.conv_input = spconv.SparseSequential(
spconv.SubMConv3d(input_channels,
16,
3,
padding=1,
bias=False,
indice_key='subm1'),
norm_fn(16),
nn.ReLU(),
)
self.conv1 = spconv.SparseSequential(
spconv.SparseSequential(
spconv.SubMConv3d(16,
16,
3,
padding=1,
bias=False,
indice_key='subm1'),
norm_fn(16),
nn.ReLU(),
))
self.conv2 = spconv.SparseSequential(
# [1600, 1408, 41] <- [800, 704, 21]
spconv.SparseSequential(
spconv.SparseConv3d(16,
32,
3,
stride=2,
padding=1,
bias=False,
indice_key='spconv2'),
norm_fn(32),
nn.ReLU(),
),
spconv.SparseSequential(
spconv.SubMConv3d(32,
32,
3,
padding=1,
bias=False,
indice_key='subm2'),
norm_fn(32),
nn.ReLU(),
),
spconv.SparseSequential(
spconv.SubMConv3d(32,
32,
3,
padding=1,
bias=False,
indice_key='subm2'),
norm_fn(32),
nn.ReLU(),
))
self.conv3 = spconv.SparseSequential(
# [800, 704, 21] <- [400, 352, 11]
spconv.SparseSequential(
spconv.SparseConv3d(32,
64,
3,
stride=2,
padding=1,
bias=False,
indice_key='spconv3'),
norm_fn(64),
nn.ReLU(),
),
spconv.SparseSequential(
spconv.SubMConv3d(64,
64,
3,
padding=1,
bias=False,
indice_key='subm3'),
norm_fn(64),
nn.ReLU(),
),
spconv.SparseSequential(
spconv.SubMConv3d(64,
64,
3,
padding=1,
bias=False,
indice_key='subm3'),
norm_fn(64),
nn.ReLU(),
))
self.conv4 = spconv.SparseSequential(
# [400, 352, 11] <- [200, 176, 5]
spconv.SparseSequential(
spconv.SparseConv3d(64,
64,
3,
stride=2,
padding=(0, 1, 1),
bias=False,
indice_key='spconv4'),
norm_fn(64),
nn.ReLU(),
),
spconv.SparseSequential(
spconv.SubMConv3d(64,
64,
3,
bias=False,
indice_key='subm4',
padding=1),
norm_fn(64),
nn.ReLU(),
),
spconv.SparseSequential(
spconv.SubMConv3d(64, 64, 3, bias=False, indice_key='subm4'),
norm_fn(64),
nn.ReLU(),
))
# 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=0, # last_pad
bias=False,
indice_key='spconv_down2'),
norm_fn(128),
nn.ReLU(),
)
#######################################接下来是RPN head #####################################
BatchNorm2d = partial(nn.BatchNorm2d, eps=1e-3, momentum=0.01)
Conv2d = partial(
nn.Conv2d, bias=False
) ########################################voxel_sa##########################################
ConvTranspose2d = partial(nn.ConvTranspose2d, bias=False)
self.block1 = nn.Sequential(
nn.ZeroPad2d(1),
Conv2d(256, 128, 3, stride=1),
BatchNorm2d(128),
nn.ReLU(),
Conv2d(128, 128, 3, padding=1), #1
BatchNorm2d(128),
nn.ReLU(),
Conv2d(128, 128, 3, padding=1), #2
BatchNorm2d(128),
nn.ReLU(),
Conv2d(128, 128, 3, padding=1), #3
BatchNorm2d(128),
nn.ReLU(),
Conv2d(128, 128, 3, padding=1), #4
BatchNorm2d(128),
nn.ReLU(),
Conv2d(128, 128, 3, padding=1), #5
BatchNorm2d(128),
nn.ReLU(),
)
self.block2 = nn.Sequential(
nn.ZeroPad2d(1),
Conv2d(128, 256, 3, stride=2),
BatchNorm2d(256),
nn.ReLU(),
Conv2d(256, 256, 3, padding=1), #1
BatchNorm2d(256),
nn.ReLU(),
Conv2d(256, 256, 3, padding=1), #2
BatchNorm2d(256),
nn.ReLU(),
Conv2d(256, 256, 3, padding=1), #3
BatchNorm2d(256),
nn.ReLU(),
Conv2d(256, 256, 3, padding=1), #4
BatchNorm2d(256),
nn.ReLU(),
Conv2d(256, 256, 3, padding=1), #5
BatchNorm2d(256),
nn.ReLU(),
)
self.deblock1 = nn.Sequential(ConvTranspose2d(128, 256, 1, stride=1),
BatchNorm2d(256), nn.ReLU())
self.deblock2 = nn.Sequential(ConvTranspose2d(256, 256, 2, stride=2),
BatchNorm2d(256), nn.ReLU())
self.rpn_head_conv_cls = nn.Conv2d(512, 6 * 3, 1)
self.rpn_head_conv_box = nn.Conv2d(512, 6 * 7, 1)
self.rpn_head_conv_dir_cls = nn.Conv2d(512, 6 * 2, 1)
#######################################接下来是voxeel SA #####################################
mlps = [[16, 16, 16], [16, 16, 16]]
self.conv1_sa_layer = pointnet2_stack_modules.StackSAModuleMSG(
radii=[0.4, 0.8],
nsamples=[16, 16],
mlps=mlps,
use_xyz=True,
pool_method='max_pool',
)
self.conv2_sa_layer = pointnet2_stack_modules.StackSAModuleMSG(
radii=[0.8, 1.2],
nsamples=[16, 32],
mlps=[[32, 32, 32], [32, 32, 32]],
use_xyz=True,
pool_method='max_pool',
)
self.conv3_sa_layer = pointnet2_stack_modules.StackSAModuleMSG(
radii=[1.2, 2.4],
nsamples=[16, 32],
mlps=[[64, 64, 64], [64, 64, 64]],
use_xyz=True,
pool_method='max_pool',
)
self.conv4_sa_layer = pointnet2_stack_modules.StackSAModuleMSG(
radii=[2.4, 4.8],
nsamples=[16, 32],
mlps=[[64, 64, 64], [64, 64, 64]],
use_xyz=True,
pool_method='max_pool',
)
c_in = 16 + 16 + 32 + 32 + 64 + 64 + 64 + 64 + 128 * 2 # C*D
self.raw_sa_layer = pointnet2_stack_modules.StackSAModuleMSG(
radii=[0.4, 0.8],
nsamples=[16, 16],
mlps=[[1, 16, 16], [1, 16, 16]],
use_xyz=True,
pool_method='max_pool',
)
c_in += 16 + 16
self.vsa_point_feature_fusion = nn.Sequential(
nn.Linear(c_in, 128, bias=False), # c_in=640
nn.BatchNorm1d(128),
nn.ReLU(),
)
self.num_point_features = 128
self.num_point_features_before_fusion = c_in
#############################################Point Head############################################
self.point_head_cls_layer = nn.Sequential(
nn.Linear(c_in, 256, bias=False), #1
nn.BatchNorm1d(256),
nn.ReLU(),
nn.Linear(256, 256, bias=False), #2
nn.BatchNorm1d(256),
nn.ReLU(),
nn.Linear(256, 1, bias=True), #3
)
###########################################RCNN########################################################
self.roi_pool_layer = pointnet2_stack_modules.StackSAModuleMSG(
radii=[0.8, 1.6],
nsamples=[16, 16],
mlps=[[128, 64, 64], [128, 64, 64]],
use_xyz=True,
pool_method="max_pool")
self.rcnn_share_layer = nn.Sequential(
nn.Conv1d(6**3 * (64 + 64), 256, kernel_size=1, bias=False),
nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(0.3),
nn.Conv1d(256, 256, kernel_size=1, bias=False),
nn.BatchNorm1d(256), nn.ReLU())
self.rcnn_cls_layer = nn.Sequential(
nn.Conv1d(256, 256, kernel_size=1, bias=False),
nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(0.3),
nn.Conv1d(256, 256, 1, bias=False), nn.BatchNorm1d(256), nn.ReLU(),
nn.Conv1d(256, 1, 1, stride=1, bias=True))
self.rcnn_reg_layer = nn.Sequential(
nn.Conv1d(256, 256, 1, stride=1, bias=False), nn.BatchNorm1d(256),
nn.ReLU(), nn.Dropout(0.3),
nn.Conv1d(256, 256, 1, stride=1, bias=False), nn.BatchNorm1d(256),
nn.ReLU(), nn.Conv1d(256, 7, 1, stride=1, bias=True))
self.voxel_size = [0.05, 0.05, 0.1]
self.point_cloud_range = [0, -40.0, -3.0, 70.4, 40.0, 1.0]
self.anchors = self.get_anchor()