简析SA-SSD使用mmdect框架生成训练数据的代码细节
1. 前言
我的上一篇博客简要分析了SA-SSD在预处理训练评估的框架。这篇博客将更近一步去分析SA-SSD使用mmdetection生成训练集的过程。考虑到SSD是基于mmdetection实现的，所以这篇博客也会分析涉及到mmdetection的一部分代码。因为我是小白，所以我会把自己不懂的地方都会记下来。
2. 简析mmdetection的训练流程
2.1 训练总体流程
考虑到SSD是基于mmdetection实现的，首先应该认识mmdetection框架下的训练流程。SSD的训练学习代码如下所示：
model = build_detector(
cfg.model, train_cfg=cfg.train_cfg, test_cfg=cfg.test_cfg)
train_dataset = get_dataset(cfg.data.train)
train_detector(
model,
train_dataset,
cfg,
distributed=distributed,
validate=args.validate,
logger=logger)
从上述代码中，可见mmdetection的训练流程大致分为三步：
第一步，初始化Detector，对应函数build_detector；
第二步，加载训练数据集，对应函数get_dataset；
第三步，训练Detector，对应函数train_detector；
在上述代码中，cfg起着非常重要的作用。cfg.model代表模型中的超参数。cfg.data.train代表训练数据集的信息。cfg记录着优化器，学习率等等训练细节。入门的小白可以仔细查看一下SA-SSD中的cfg文件。
函数train_detector是一个比较自动的模块。它会根据cfg中设定的训练细节训练Detector。我不需要修改它，因此train_detector就不是我的重点。作为小白，我比较关注get_dataset和build_detector。在SA-SSD中，训练数据只是视场范围内的点云。假设我的模型需要雷达全部点云呢，或是结合RGB图像呢，或是结合双目RGB图像呢，抑或考虑IMU信息呢，我怎样去生成我所需要的train_dataset呢？这是一个很大的问题。build_detector的重要性不言而喻。这篇博客将分析第一个问题。
2.2 深入理解get_dataset
首先理解cfg.data.train。在上一篇博客已经贴出它的信息，这里就不再重复。在cfg.data.train中，与训练数据相关的参数是：数据集路径，图像尺寸，图像裁剪比（跟FPN网络相关），图像归一化参数，检测目标（比如指定只检测汽车的3D目标），3D目标真值，点云体素化参数，Anchor生成参数。
再去理解函数get_dataset。上一篇文章谈到get_dataset代码中多次出现obj_from_dict。对这个核心函数的分析可见这篇博客。代码如下所示：
def obj_from_dict(info, parrent=None, default_args=None):
"""Initialize an object from dict.
The dict must contain the key "type", which indicates the object type, it
can be either a string or type, such as "list" or ``list``. Remaining
fields are treated as the arguments for constructing the object.
Args:
info (dict): Object types and arguments.
module (:class:`module`): Module which may containing expected object
classes.
default_args (dict, optional): Default arguments for initializing the
object.
Returns:
any type: Object built from the dict.
"""
assert isinstance(info, dict) and 'type' in info
assert isinstance(default_args, dict) or default_args is None
args = info.copy()
obj_type = args.pop('type')
if mmcv.is_str(obj_type):
if parrent is not None:
obj_type = getattr(parrent, obj_type)
else:
obj_type = sys.modules[obj_type]
elif not isinstance(obj_type, type):
raise TypeError('type must be a str or valid type, but got {}'.format(
type(obj_type)))
if default_args is not None:
for name, value in default_args.items():
args.setdefault(name, value)
return obj_type(**args)
看上去挺复杂的。注释说这段函数的作用是Initialize an object from dict，通俗理解是根据字典型变量info去指定初始化一个parrent类对象。说白了，就是字典型变量中储存了类的初始化变量。核心调用是getattr。总之，obj_from_dict是一种做指定初始化的功能函数。
Ok，理解了obj_from_dict，是时候刚一波get_dataset的源码：
def get_dataset(data_cfg):
# 生成index文件的实例，'ann_file'是data_root + 'ImageSets/train.txt'
# num_dset 就是训练数据总数
if isinstance(data_cfg['ann_file'], (list, tuple)):
ann_files = data_cfg['ann_file']
num_dset = len(ann_files)
else:
ann_files = [data_cfg['ann_file']]
num_dset = 1
# SA-SSD没有使用它，按照else，生成 N 个 None
if 'proposal_file' in data_cfg.keys():
if isinstance(data_cfg['proposal_file'], (list, tuple)):
proposal_files = data_cfg['proposal_file']
else:
proposal_files = [data_cfg['proposal_file']]
else:
proposal_files = [None] * num_dset
assert len(proposal_files) == num_dset
# SA-SSD没有使用它，算法不需要图像，'img_prefix'=None
# 按照else，生成 N 个 None
# 如果需要RGB的话，可以在cfg中写img_prefix=data_root + 'train2017/'相应路径
if isinstance(data_cfg['img_prefix'], (list, tuple)):
img_prefixes = data_cfg['img_prefix']
else:
img_prefixes = [data_cfg['img_prefix']] * num_dset
assert len(img_prefixes) == num_dset
# 按照data_cfg['generator']的参数，初始化voxel_generator，用于预处理点云体素化
if 'generator' in data_cfg.keys() and data_cfg['generator'] is not None:
generator = obj_from_dict(data_cfg['generator'], voxel_generator)
else:
generator = None
# 按照data_cfg['augmentor']的参数，初始化point_augmentor，用于提供3D目标真值
if 'augmentor' in data_cfg.keys() and data_cfg['augmentor'] is not None:
augmentor = obj_from_dict(data_cfg['augmentor'], point_augmentor)
else:
augmentor = None
# 按照data_cfg['anchor_generator']的参数，初始化anchor3d_generator，用于提供3DAnchor
if 'anchor_generator' in data_cfg.keys() and data_cfg['anchor_generator'] is not None:
anchor_generator = obj_from_dict(data_cfg['anchor_generator'], anchor3d_generator)
else:
anchor_generator = None
# 按照data_cfg['target_encoder']的参数，初始化bbox3d_target
# SA-SSD中貌似没有使用，返回 None
if 'target_encoder' in data_cfg.keys() and data_cfg['target_encoder'] is not None:
target_encoder = obj_from_dict(data_cfg['target_encoder'], bbox3d_target)
else:
target_encoder = None
dsets = []
# 装填用于训练的数据
for i in range(num_dset):
# 定义字典型变量data_info ，用于引导训练数据的装填
data_info = copy.deepcopy(data_cfg)
data_info['ann_file'] = ann_files[i]
data_info['proposal_file'] = proposal_files[i]
data_info['img_prefix'] = img_prefixes[i]
if generator is not None:
data_info['generator'] = generator
if anchor_generator is not None:
data_info['anchor_generator'] = anchor_generator
if augmentor is not None:
data_info['augmentor'] = augmentor
if target_encoder is not None:
data_info['target_encoder'] = target_encoder
# 使用data_info去实例化datasets
dset = obj_from_dict(data_info, datasets)
dsets.append(dset)
if len(dsets) > 1:
# 从上述操作中，每一个训练数据都是一个datasets类
# 使用ConcatDataset，把所有datasets类，统一变成一类datasets类
dset = ConcatDataset(dsets)
else:
dset = dsets[0]
return dset
上述代码中，最为核心的是这句话dset = obj_from_dict(data_info, datasets)。我需要看看类datasets类的实例化过程。
然而datasets类是一个Virtual类的存在，
__all__ = [
'CustomDataset', 'CocoDataset', 'GroupSampler', 'DistributedGroupSampler',
'ConcatDataset', 'build_dataloader', 'to_tensor', 'random_scale',
'show_ann', 'get_dataset', 'KittiLiDAR','KittiVideo', 'VOCDataset'
]
回头再想想obj_from_dict，更有深层次的理解。根据字典型变量info去指定初始化一个parrent类对象。如果parrent类是一个虚类，它会根据info的变量自动地匹配一个Matched的子类，去指定初始化这个子类的实例。
毫无疑问，肯定是生成datasets类子类中的KittiLiDAR类。它会在下一节分析。
2.3 简析KittiLiDAR类
该类的初始化代码如下所示。从下面代码可以看出字典型变量data_info的内部变量跟KittiLiDAR初始化所需要的变量是匹配的。
class KittiLiDAR(Dataset):
def __init__(self, root, ann_file,
img_prefix,
img_norm_cfg,
img_scale=(1242, 375),
size_divisor=32,
proposal_file=None,
flip_ratio=0.5,
with_point=False,
with_mask=False,
with_label=True,
class_names = ['Car', 'Van'],
augmentor=None,
generator=None,
anchor_generator=None,
anchor_area_threshold=1,
target_encoder=None,
out_size_factor=2,
test_mode=False):
self.root = root
self.img_scales = img_scale if isinstance(img_scale,
list) else [img_scale]
assert mmcv.is_list_of(self.img_scales, tuple)
# normalization configs
self.img_norm_cfg = img_norm_cfg
# flip ratio
self.flip_ratio = flip_ratio
# size_divisor (used for FPN)
self.size_divisor = size_divisor
self.class_names = class_names
self.test_mode = test_mode
self.with_label = with_label
self.with_mask = with_mask
self.with_point = with_point
# 获取KITTI相关各种数据的前缀路径
self.img_prefix = osp.join(root, 'image_2')
self.right_prefix = osp.join(root, 'image_3')
self.lidar_prefix = osp.join(root, 'velodyne_reduced')
self.calib_prefix = osp.join(root, 'calib')
self.label_prefix = osp.join(root, 'label_2')
with open(ann_file, 'r') as f:
self.sample_ids = list(map(int, f.read().splitlines()))
if not self.test_mode:
self._set_group_flag()
# transforms
self.img_transform = ImageTransform(
size_divisor=self.size_divisor, **self.img_norm_cfg)
# voxel
self.augmentor = augmentor
self.generator = generator
self.target_encoder = target_encoder
self.out_size_factor = out_size_factor
self.anchor_area_threshold = anchor_area_threshold
# anchor
if anchor_generator is not None:
feature_map_size = self.generator.grid_size[:2] // self.out_size_factor
feature_map_size = [*feature_map_size, 1][::-1]
anchors = anchor_generator(feature_map_size)
self.anchors = anchors.reshape([-1, 7])
self.anchors_bv = rbbox2d_to_near_bbox(
self.anchors[:, [0, 1, 3, 4, 6]])
else:
self.anchors=None
DataLoader中，会需要这个类的__getitem__函数：
def __getitem__(self, idx):
if self.test_mode:
return self.prepare_test_img(idx)
while True:
data = self.prepare_train_img(idx)
if data is None:
idx = self._rand_another(idx)
continue
return data
来看看函数prepare_train_img，输出字典型变量data。
def prepare_train_img(self, idx):
sample_id = self.sample_ids[idx]
# load image
img = mmcv.imread(osp.join(self.img_prefix, '%06d.png' % sample_id))
img, img_shape, pad_shape, scale_factor = self.img_transform(img, 1, False)
objects = read_label(osp.join(self.label_prefix, '%06d.txt' % sample_id))
calib = Calibration(osp.join(self.calib_prefix, '%06d.txt' % sample_id))
gt_bboxes = [object.box3d for object in objects if object.type not in ["DontCare"]]
gt_bboxes = np.array(gt_bboxes, dtype=np.float32)
gt_types = [object.type for object in objects if object.type not in ["DontCare"]]
#gt_labels = np.ones(len(gt_bboxes), dtype=np.int64)
# transfer from cam to lidar coordinates
if len(gt_bboxes) != 0:
gt_bboxes[:, :3] = project_rect_to_velo(gt_bboxes[:, :3], calib)
img_meta = dict(
img_shape=img_shape,
sample_idx=sample_id,
calib=calib
)
data = dict(
img=to_tensor(img),
img_meta = DC(img_meta, cpu_only=True)
)
if self.anchors is not None:
data['anchors'] = DC(to_tensor(self.anchors.astype(np.float32)))
if self.with_mask:
NotImplemented
if self.with_point:
points = read_lidar(osp.join(self.lidar_prefix, '%06d.bin' % sample_id))
if self.augmentor is not None and self.test_mode is False:
sampled_gt_boxes, sampled_gt_types, sampled_points = self.augmentor.sample_all(gt_bboxes, gt_types)
assert sampled_points.dtype == np.float32
gt_bboxes = np.concatenate([gt_bboxes, sampled_gt_boxes])
gt_types = gt_types + sampled_gt_types
assert len(gt_types) == len(gt_bboxes)
# to avoid overlapping point (option)
masks = points_in_rbbox(points, sampled_gt_boxes)
points = points[np.logical_not(masks.any(-1))]
# paste sampled points to the scene
points = np.concatenate([sampled_points, points], axis=0)
# select the interest classes
selected = [i for i in range(len(gt_types)) if gt_types[i] in self.class_names]
gt_bboxes = gt_bboxes[selected, :]
gt_types = [gt_types[i] for i in range(len(gt_types)) if gt_types[i] in self.class_names]
# force van to have same label as car
gt_types = ['Car' if n == 'Van' else n for n in gt_types]
gt_labels = np.array([self.class_names.index(n) + 1 for n in gt_types], dtype=np.int64)
self.augmentor.noise_per_object_(gt_bboxes, points, num_try=100)
gt_bboxes, points = self.augmentor.random_flip(gt_bboxes, points)
gt_bboxes, points = self.augmentor.global_rotation(gt_bboxes, points)
gt_bboxes, points = self.augmentor.global_scaling(gt_bboxes, points)
if isinstance(self.generator, VoxelGenerator):
#voxels, coordinates, num_points = self.generator.generate(points)
voxel_size = self.generator.voxel_size
pc_range = self.generator.point_cloud_range
grid_size = self.generator.grid_size
keep = points_op_cpu.points_bound_kernel(points, pc_range[:3], pc_range[3:])
voxels = points[keep, :]
coordinates = ((voxels[:, [2, 1, 0]] - np.array(pc_range[[2,1,0]], dtype=np.float32)) / np.array(
voxel_size[::-1], dtype=np.float32)).astype(np.int32)
num_points = np.ones(len(keep)).astype(np.int32)
data['voxels'] = DC(to_tensor(voxels.astype(np.float32)))
data['coordinates'] = DC(to_tensor(coordinates))
data['num_points'] = DC(to_tensor(num_points))
if self.anchor_area_threshold >= 0 and self.anchors is not None:
dense_voxel_map = sparse_sum_for_anchors_mask(
coordinates, tuple(grid_size[::-1][1:]))
dense_voxel_map = dense_voxel_map.cumsum(0)
dense_voxel_map = dense_voxel_map.cumsum(1)
anchors_area = fused_get_anchors_area(
dense_voxel_map, self.anchors_bv, voxel_size, pc_range, grid_size)
anchors_mask = anchors_area > self.anchor_area_threshold
data['anchors_mask'] = DC(to_tensor(anchors_mask.astype(np.uint8)))
# filter gt_bbox out of range
bv_range = self.generator.point_cloud_range[[0, 1, 3, 4]]
mask = filter_gt_box_outside_range(gt_bboxes, bv_range)
gt_bboxes = gt_bboxes[mask]
gt_labels = gt_labels[mask]
else:
NotImplementedError
# skip the image if there is no valid gt bbox
if len(gt_bboxes) == 0:
return None
# limit rad to [-pi, pi]
gt_bboxes[:, 6] = limit_period(
gt_bboxes[:, 6], offset=0.5, period=2 * np.pi)
if self.with_label:
data['gt_labels'] = DC(to_tensor(gt_labels))
data['gt_bboxes'] = DC(to_tensor(gt_bboxes))
return data
字典型变量data包含img,img_meta,anchors,voxel(视场范围的点云),coordinates,num_points,anchor_mask,gt_labes,gt_bboxes。总之它包含输入值也包含3D目标检测真值。
函数prepare_test_img构造和prepare_train_img几乎一样。
2.4 理解训练过程的一个问题
这一节分析2.1节提及的函数train_detector。我一直都有一个疑问，就是train_dataset提供那么多数据，它是怎样把指定的变量喂给Model的？比如SA-SSD只需要喂点云，而train_dataset有那么多数据。代码在哪里指定输入对象的？
# 前向计算过程中，是直接把data塞进model中的
# 从KITTILiDAR可知，data包含img，img_meta，anchors，voxel等等变量
# 在前向计算中，是哪些变量发挥作用呢？
for data in enumerate(data_loader):
results = model(**data)
找了半天，答案并不在函数train_detector。
答案在函数build_detector中。SA-SSD属于父类SingleStageDetector。父类中前向计算函数forward就指明了模型的输入：
# img, img_meta, **kwargs都属于train_dataset push的数据
def forward_train(self, img, img_meta, **kwargs):
batch_size = len(img_meta)
ret = self.merge_second_batch(kwargs)
vx = self.backbone(ret['voxels'], ret['num_points'])
# 后续代码略去
函数merge_second_batch代码如下所示。kwargs是字典型变量，从它那里搜索voxels和num_points作为输入。3D检测真值gt_labels和gt_bboxes也将保存。
def merge_second_batch(self, batch_args):
ret = {}
for key, elems in batch_args.items():
if key in [
'voxels', 'num_points',
]:
ret[key] = torch.cat(elems, dim=0)
elif key == 'coordinates':
coors = []
for i, coor in enumerate(elems):
coor_pad = F.pad(
coor, [1, 0, 0, 0],
mode='constant',
value=i)
coors.append(coor_pad)
ret[key] = torch.cat(coors, dim=0)
elif key in [
'img_meta', 'gt_labels', 'gt_bboxes',
]:
ret[key] = elems
else:
ret[key] = torch.stack(elems, dim=0)
return ret
2.5 阶段性小结
对于SA-SSD来说，2.1节提及的train_dataset是由全体训练数据以KittiLiDAR类构成的。
3. 小结
mmdetection是一个工程性很强的开源框架。代码稍微抽象，封装性强。相比上一篇博客，这篇博客更加细致地分析了输入数据的流向。可能会有别人说，直接使用它难道不香吗。博主作为小白，在使用mmdetection框架遇到bug，只有明白框架的来龙去脉，才能有能力去修改它。下一篇博客将会分析SA-SSD的网络架构。