跳过正文
  1. 技术博客/

Raft训练日志2

·1196 字·3 分钟·
Somnus
作者
Somnus
记录学习、项目、实验和一些被密码保护起来的日常。

前情提要
#

我上一个训练日志中,尝试对RAFT-Stereo模型进行优化,通过添加refinement层和边缘损失感知,然后训练完的模型的EPE和D1几乎没有产生变化。因为这个层结构对整个模型来说都没什么影响。接着我在改完的这个模型后,对模型进行了partial unfreezement,然后在此基础上对middlebury数据集进行训练。最后得到的结果EPE和D1反而变差了,这个大概率是因为我本地的数据集太小,在此基础上进行训练会很容易导致过拟合的现象。自此,我便将研究重心转移到模型轻量化方面来。


模型设计
#

首先,我先去官方库上找到了不同版本模型的下载链接。我原本一直都是使用他的middlebury版本,因为我本地的数据集也是middelbury,得到的效果也很好,然后我注意到除了middlebury,以外还有两个以数据集命名的预训练模型eth3d和sceneflow。还有一个模型就是raftstereo-realtime.pt,我后来调查得到,这个模型就是官方训练出来的轻量化版本,然,我就想着先拿这个模型测试一下。 除此以外,我还想了通过转换精度来实现轻量化,具体就是不改变模型原始网络结构,也不用重新训练模型,将原本的Float32转换成Float16,这样便可以大幅度降低模型的大小和显存占用,对于低显存的边缘设备,能提升帧率和处理速度。以下为处理代码,处理逻辑实际上value.half()这一句。

 1def convert_to_fp16(input_path, output_path):
 2    print(f"Reading original checkpoint from: {input_path}")
 3    state_dict = torch.load(input_path, map_location='cpu')
 4    new_state_dict = {}
 5    for key, value in state_dict.items():
 6        if isinstance(value, torch.Tensor) and value.is_floating_point():
 7            new_state_dict[key] = value.half()  # Convert to float16
 8        else:
 9            new_state_dict[key] = value  # Keep non-floating-point parameters unchanged
10    print(f"Saving FP16 checkpoint to: {output_path}")
11    torch.save(new_state_dict, output_path)
12    print("Conversion completed successfully!")

这一方案的核心思想就是:

  • 保留原始高精度模型结构。
  • 避免重新训练带来的额外成本。

训练结果
#

官方 realtime 模型与高精度模型对比
#

模型参数量EPED1
RAFT-Stereo Middlebury11.14M2.375812.0526
RAFT-Stereo Realtime9.89M11.155728.5460

从结果可以看出,官方 realtime 模型参数量相较高精度 Middlebury 模型有所减少,但在 Middlebury2014 数据集上的精度下降非常明显:

  • EPE 由 2.3758 上升到 11.1557
  • D1 由 12.0526 上升到 28.5460

这说明:直接采用官方轻量模型虽然能够降低模型规模,但会带来较大的精度损失,不适合作为当前任务的最终部署方案。

FP16 半精度部署结果
#

模型EPED1平均推理时间FPS最大显存占用模型大小
Baseline FP322.375812.0526209.80 ms4.77844.91 MB43 MB
FP162.353512.0593173.59 ms5.76662.34 MB22 MB

进一步计算可得:

  • 模型大小:43 MB → 22 MB,下降约 48.8%
  • 显存占用:844.91 MB → 662.34 MB,下降约 21.6%
  • 平均推理时间:209.80 ms → 173.59 ms,缩短约 17.3%
  • FPS:4.77 → 5.76,提升约 20.8%

而在精度方面:

  • EPE:2.3758 → 2.3535,基本保持一致
  • D1:12.0526 → 12.0593,几乎不变

这说明:将高精度 Middlebury 模型直接转换为 FP16 半精度模型后,模型精度基本保持稳定,同时明显降低了模型存储、显存占用和推理时间。

小结
#

虽然我做出的改进非常小,但是得到的效果我觉得还是可以的,应该是足够支撑去作为本科毕业设计的成果的,但不知道胡老师会怎么看,希望能美美通关。