深度估计自监督模型monodepth2在自己数据集的实战——单卡/多卡训练、推理、Onnx转换和量化指标评估

本文详细介绍monodepth2模型在自己数据集的实战方法，包括单卡/多卡训练、推理、Onnx转换和量化评估等，关于理论部分请参见另一篇博客：深度估计自监督模型monodepth2论文总结和源码分析

一、训练过程：

（一）训练过程：

1.数据集准备：一个是要有连续帧数据（前后帧图片编号相差1）1-3W张左右，相机内参矩阵，还有可选的左右视图数据和点云代表的深度真值。

2.按照源码KITTI数据集格式生成train_files.txt和val_files.txt

3.继承MonoDataset类编写自己的数据类，主要内容是设置内参、分辨率、数据读取等

4.把option.py的dataset和split参数设置为自己的数据类，然后在trainer.py中对应到自己数据集的地址

都准备好以后直接指定模型类型就可以训练了,命令为python train.py --model_name 训练开始时，会在trainer类中初始化所有网络、data_loader，然后调用train()函数分batch训练。

（二）多卡训练：

源码只有单卡训练代码，想要实现多卡训练可以用torch.nn.parallel.DistributedDataParallel（DDP）进行改造，但比较繁琐，一是因为整个模型由四个子网络拼接而成，需要一个个设置，二是源码中很多tensor的第一维都默认为batch_size，多卡情况下假如最后一个batch的数量不足一个batch_size,这个维度的大小会报错，需要手动改成可配置的大小。

步骤主要包括：

1.参数设置和初始化，可用torch.multiprocessing.spawn开启多进程，使用cuda号作为进程编号

dist.init_process_group("nccl", rank=rank, world_size=world_size)
if opts.ddp:
   world_size = opts.world_size
   torch.multiprocessing.spawn(main_worker,
                               args=(world_size,),
                               nprocs=world_size,
                               join=True)

2.把每个子网络都用DistributedDataParallel包装起来，放到对应的cuda上，并使用SyncBatchNorm.convert_sync_batchnorm实现全局同步BN

if self.opt.ddp:
    self.models["depth"] = SyncBatchNorm.convert_sync_batchnorm(self.models["depth"])
    self.models["depth"].to(self.rank)
    self.models["depth"] = DistributedDataParallel(self.models["depth"], device_ids=[rank],
                                                   find_unused_parameters=True)

3.把训练集和验证集都用torch.utils.data.distributed.DistributedSampler包装起来，实现对数据集的多卡分配，作为sampler传入data_loader

if self.opt.ddp:
    self.train_sampler = DistributedSampler(train_dataset)
    self.train_loader = DataLoader(
    train_dataset, self.opt.batch_size, sampler=self.train_sampler, num_workers=2)

4.训练的时候要为sampler设置epoch值作为随机种子，否则数据无法被shuffle

for self.epoch in range(self.opt.num_epochs):
    if self.opt.ddp:
       self.train_sampler.set_epoch(self.epoch)

5.损失函数更新的时候要用torch.distributed.reduce进行全局更新。每个epoch结束时需要全局平均的validate损失的话要手动计算，并在0号cuda输出

def reduce_loss(self, tensor, rank, world_size):
    with torch.no_grad():
        dist.reduce(tensor, dst=0)
        if rank == 0:
            tensor /= world_size
    return tensor

6.多卡模型保存时候要用save = model.module.state_dict()，载入的时候要用pretrained_dict = torch.load(path, map_location=lambda storage, loc: storage.cuda(self.rank))实现cpu到gpu的转换。

这里面还有一个问题，就是源码很多地方输入的第一维默认为batchsize，改成多卡以后DistributedSampler并不能保证最后一个batch的大小是batch_size,这里要手动调整一下输入大小。

（三）tensorboard可视化

源码中有tensorboard相关代码，会默认将训练、验证相关参数和图片保存在_{/tmp/models/train和}/tmp/models/val路径，然后可以用tensorboard --logdir 命令来查看，关于tensorboard相关使用方法可以参考以下博客：苹果姐：pytorch使用tensorboard实现可视化总结

二、推理：

源码提供了test_simple代码，内容首先如果目录中没有模型文件则先下载模型并解压成五个pth文件，分别对应到四个网络，用torch.load加载各自的参数进行实例化，然后设置为eval模式。如果要用自己训练的模型，则需要在代码中手动指定模型地址，默认地址是~/tmp/models/mode_name

源码中的推理过程是：对输入图片用pil工具resize成模型需要的大小后先输入encoder，再输入decoder，得到的结果中取出第0个尺度的disp图，上采样到图片原始分辨率，然后按照按照0.1到100的尺度将disp图取倒数得到深度图，再配上不同的颜色输出保存。

源码可以支持对单张图片和整个文件夹的推理，结果保存’disp.jpg’格式的可视化深度图文件和’_disp.npy’格式的numpy数据文件。

三、转onnx并验证

源码未提供现成代码，但部分可参考test_simple.py。过程是下载并初始化模型后自定义一个完整的模型，把encoder和decoder合并，然后产生与输入图片尺度相同的随机tensor输入网络中，再用torch.onnx.export()函数导出onnx。这里面要注意的问题是这个函数不支持比较复杂的输出，如字典等，而原模型的输出是一个很大的字典，所以在自定义模型中只输出了需要的tensor。最好还要定义输入输出的名字，对于多输出输出需要定义名字列表，以备推理使用。核心代码如下：

class FullModel(nn.Module):
    def __init__(self, encoder, decoder):
        super(FullModel, self).__init__()
        self.encoder = encoder
        self.decoder = decoder

    def forward(self, x):
        features = self.encoder(x)
        outputs = self.decoder(features)
        return outputs[("disp", 0)]

full_model = FullModel(encoder, depth_decoder)
full_model = full_model.eval()

x = torch.randn((1,3,192,640), device="cuda")
output = full_model(x)
output_path = "outputs/monodepth2.onnx"

torch.onnx.export(full_model,               # model being run
                   x,                       # model input (or a tuple for multiple inputs)
                  output_path,   # where to save the model (can be a file or file-like object)
                  export_params=True,        # store the trained parameter weights inside the model file
                  opset_version=11,         # the ONNX version to export the model to
                  do_constant_folding=True,  # whether to execute constant folding for optimization
                  input_names = ['inputx'],   # the model's input names
                  output_names = ['outputy'], # the model's output names
                  verbose=True
                  )

转成功以后需要进行推理测试，也有点麻烦，因为要用到onnxruntime框架，他要求的输入输出都是numpy格式，要先将输入图像升一维并转成numpy,再调用onnxruntime.InferenceSession(onnx_path).run(ouput_name,input_feed={})执行推理，input_feed需要把每个输入和网络的输入名字对应上。得出结果转回tensor,执行和test_sample一样的操作，发现结果相同即搞定。具体可参考：onnx模型推理(python) - osc_kwludyt3的个人空间 - OSCHINA - 中文开源技术交流社区

四、模型评估

源码中还提供了点云转深度图、深度图评估代码，可以计算"abs_rel", “sq_rel”, “rmse”, “rmse_log”, “a1”, “a2”, "a3"等指标，点云转深度图的地方有一部分代码是用来去重的：

# find the duplicate points and choose the closest depth
inds = sub2ind(depth.shape, velo_pts_im[:, 1], velo_pts_im[:, 0])
dupe_inds = [item for item, count in Counter(inds).items() if count > 1]
for dd in dupe_inds:
    pts = np.where(inds == dd)[0]
    x_loc = int(velo_pts_im[pts[0], 0])
    y_loc = int(velo_pts_im[pts[0], 1])
    depth[y_loc, x_loc] = velo_pts_im[pts, 2].min()
depth[depth < 0] = 0

意思是假如两个点投影到同一个像素点，只保留深度较近的，这段第一次看有点费解，其实前面就是对每个点按照行列的索引来联合排序，找到重复的点。还有一个地方我认为可以优化：

源码在去重以后再过滤掉深度<0的点，但这时候可能在去重的时候假如同一个像素点深度有正有负，则只保留了深度为负数的点，因为深度为负数的肯定比正数要小，这样可能误删了很多点。所以我认为应该先过滤掉深度为负数的点，再进行去重。

例如在这段代码：

# check if in bound

val_inds = (velo_pts_im[:, 0] >= 0) & (velo_pts_im[:, 1] >= 0) & (velo_pts_im[:, 2] > 0)
val_inds = val_inds & (velo_pts_im[:, 0] < im_shape[0]) & (velo_pts_im[:, 1] < im_shape[1])
velo_pts_im = velo_pts_im[val_inds, :]

第一行加上了& (velo_pts_im[:, 2] > 0)，而在去重的最后一行去掉了depth[depth < 0] = 0，输出的结果更加稠密。假如要将输出的点云投影稀疏深度图可视化，可以使用cv2.applyColorMap(gt_depth, cv2.COLORMAP_JET)进行上色。

评估的过程也可以自由调整，kitti源码设置了最大评测距离是80米，图像有效区域也进行了裁剪：

 crop = np.array([0.40810811 * gt_height, 0.99189189 * gt_height,
                  0.03594771 * gt_width,  0.96405229 * gt_width]).astype(np.int32)

这个可以按照需要来调整，另外最后的深度图尺度恢复，源码用的是中值恢复，也可以改成其他的策略：

if not opt.disable_median_scaling:
   ratio = np.median(gt_depth) / np.median(pred_depth)
   ratios.append(ratio)
   pred_depth *= ratio

pred_depth[pred_depth < MIN_DEPTH] = MIN_DEPTH
pred_depth[pred_depth > MAX_DEPTH] = MAX_DEPTH

errors.append(compute_errors(gt_depth, pred_depth))