PVN3D 模型训练 Bug 调试指南

1. 适用场景

本文针对当前仓库中 PVN3D 的 LINEMOD 训练流程，重点覆盖以下问题：

• 在 pvn3d-dev 容器内训练时出现 CUDA out of memory
• 调小 n_sample_points 后出现张量形状不匹配
• 训练脚本出现 lr_scheduler.step() 调用顺序警告
• 如何快速判断是显存问题、数据问题还是代码问题
• 如何在当前仓库内做最小代价的修复和重试

本文默认你的环境是：

• 容器：pvn3d-dev
• 训练环境：conda activate pvn3d
• 工作目录：/workspace/workflow/self/PVN3D/pvn3d

典型训练命令：

python -m train.train_linemod_pvn3d --cls ape

---

2. 本次问题现象

在执行：

python -m train.train_linemod_pvn3d --cls ape

后，训练开始初始化数据集和模型，但在首轮训练前向过程中报错：

RuntimeError: CUDA out of memory. Tried to allocate 114.00 MiB

调用栈显示报错位置出现在：

• pvn3d/lib/pvn3d.py
• pvn3d/lib/pspnet.py
• pvn3d/lib/extractors.py

更具体地说，是在 RGB 主干网络的卷积阶段触发的显存不足。

日志中同时还出现了两类警告：

1. lr_scheduler.step() 在 optimizer.step() 之前调用
2. scheduler.step(epoch) 这种带 epoch 参数的调用方式已被 PyTorch 标记为过时

需要区分：

• OOM 是导致训练中断的主问题
• lr_scheduler 警告不会直接导致训练崩溃，但说明训练脚本需要修正

在为了解决 OOM 而降低采样点数后，还可能继续遇到第二类问题：

RuntimeError: Sizes of tensors must match except in dimension 1.
Expected size 1 but got size 3 for tensor number 2 in the list.

这类报错出现在：

• pvn3d.py

对应位置是 DenseFusion 中对 feat_1、feat_2 和 ap_x 的拼接。

---

3. 根因分析

3.1 主因：默认配置对 7.62 GiB 显存过大

当前仓库默认配置位于：

• common.py

其中关键参数是：

• mini_batch_size = 24
• val_mini_batch_size = 24
• n_sample_points = 8192 + 4096 = 12288

这组配置对 PVN3D 来说开销很大，因为模型同时处理：

• RGB 分支特征提取
• 点云分支特征提取
• RGBD 融合
• 关键点偏移与中心偏移预测

而你的 GPU 信息显示：

• 总显存：7.62 GiB
• 已分配：5.40 GiB
• 预留：5.56 GiB
• 可用：70.44 MiB

这意味着即使只再申请 114 MiB 也会失败。

所以根因不是：

• 数据缺失
• checkpoint 损坏
• 某个 tensor 维度错乱

而是：

• 默认 batch size 和采样点数对当前显存配置过大

3.2 次因：学习率调度器调用顺序不符合当前 PyTorch 约定

训练脚本位置：

• train_linemod_pvn3d.py

原始逻辑是在每轮训练里先调用：

self.lr_scheduler.step(it)

再调用：

self.optimizer.step()

这会触发当前 PyTorch 的警告，因为推荐顺序是：

1. optimizer.step()
2. lr_scheduler.step()

这个问题本身不会造成 OOM，但会影响学习率调度的正确性，并污染训练日志。

3.3 第二类问题：DenseFusion 对采样点数写死，导致形状不匹配

报错位置在：

• pvn3d.py

旧实现中，DenseFusion 使用了固定池化：

self.ap1 = torch.nn.AvgPool1d(num_points)

然后在前向里执行：

ap_x = self.ap1(rgbd)
ap_x = ap_x.view(-1, 1024, 1).repeat(1, 1, n_pts)
return torch.cat([feat_1, feat_2, ap_x], 1)

这里隐含了一个前提：

• 模型初始化时传入的 num_points
• 数据集实际输出的点数 n_pts

必须完全一致。

一旦你为了省显存而通过命令行调低 --n_sample_points，但模型和数据集没有同步使用同一份配置，就会出现：

• feat_1 和 feat_2 的 batch / 点数维度来自真实输入
• ap_x 的形状来自写死的池化窗口
• 最终 torch.cat 时报尺寸不一致

这次报错的根因不是：

• 显存不够
• 数据缺失
• 标签错乱

而是：

• DenseFusion 对输入点数的假设过强
• LM_Dataset 没有共享训练脚本里覆盖后的 config

---

4. 如何快速判断是不是显存问题

如果你看到以下特征，基本可以直接判断是显存不足：

• 错误信息中明确出现 CUDA out of memory
• 调用栈落在 conv2d、forward、layer3 等特征提取层
• 错误里给出了 Tried to allocate xxx MiB
• GPU 总显存本来就不大，例如 8 GiB 左右
• 训练一开始就崩，而不是跑若干 epoch 后崩

相反，如果是数据问题，通常会看到：

• 文件不存在
• shape mismatch
• KeyError
• 数据路径报错
• invalid device function

这次你的日志完全符合“典型首轮前向 OOM”。

---

5. 当前仓库已做的修复

本次已修改训练脚本：

• train_linemod_pvn3d.py
• pvn3d.py
• linemod_dataset.py

5.1 新增可覆盖的训练参数

新增了以下命令行参数：

• --mini_batch_size
• --val_mini_batch_size
• --num_workers
• --n_sample_points

这样不必每次去改 common.py，可以直接在命令行临时压缩配置。

5.2 修复学习率调度器调用顺序

将训练循环调整为：

1. loss.backward()
2. optimizer.step()
3. lr_scheduler.step()

这样可以消除 PyTorch 关于顺序的警告，并避免跳过初始学习率值。

5.3 保持改动范围最小

没有改：

• 数据集逻辑
• loss 定义
• checkpoint 格式

所以这是一组低风险修复，主要用于让小显存卡也能跑起来，并避免降低点数后再触发形状错误。

5.4 修复 DenseFusion 的固定池化问题

位置：

• pvn3d.py

修复前依赖：

• AvgPool1d(num_points)

修复后改为：

• adaptive_avg_pool1d(rgbd, 1)

这意味着全局特征池化不再依赖写死的采样点数，只要输入是 B x C x N，就会稳定输出 B x C x 1。

5.5 让 LM_Dataset 复用训练脚本中的配置

位置：

• linemod_dataset.py
• train_linemod_pvn3d.py

修复前的问题是：

• 训练脚本通过命令行修改了 config.n_sample_points
• 但 LM_Dataset 内部又自己创建了一份新的 Config

结果就是：

• 模型拿到的是一套点数配置
• 数据集输出的是另一套点数配置

修复后：

• LM_Dataset 支持接收外部 config
• 训练脚本会把同一份 config 传给训练集、验证集和测试集

这样 --n_sample_points 才会真正一致地作用到模型和数据上。

---

6. 推荐的重试方式

进入容器并激活环境：

conda activate pvn3d
cd /workspace/workflow/self/PVN3D/pvn3d

6.1 第一组推荐参数

先用这组：

PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:64 python -m train.train_linemod_pvn3d \
  --cls ape \
  --mini_batch_size 2 \
  --val_mini_batch_size 2 \
  --n_sample_points 8192 \
  --num_workers 4

这组配置的思路是：

• 批大小从 24 降到 2
• 采样点数从 12288 降到 8192
• 降低数据加载并发，减少宿主机/容器压力
• 用 max_split_size_mb:64 缓解显存碎片

6.2 如果仍然 OOM

继续降到：

PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:64 python -m train.train_linemod_pvn3d \
  --cls ape \
  --mini_batch_size 1 \
  --val_mini_batch_size 1 \
  --n_sample_points 4096 \
  --num_workers 2

这组基本是“小显存保守起步配置”。

---

7. 各参数对训练的影响

7.1 mini_batch_size

作用：

• 直接决定单次前向/反向传播占用的显存

规律：

• 越大越容易 OOM
• 是最有效的降显存手段

代价：

• 训练吞吐下降
• 梯度估计更噪声

7.2 n_sample_points

作用：

• 决定每个样本保留多少点云点参与后续网络计算

规律：

• 越大显存越高
• 对点云分支和融合模块压力明显

代价：

• 降得太低可能影响姿态估计精度

实践建议：

• 先从 8192 试
• 不行再降到 4096

7.3 num_workers

作用：

• 控制 DataLoader 的并发加载进程数

规律：

• 不直接决定 GPU OOM
• 但过大时会增加 CPU 内存和容器负载

建议：

• 容器里优先设为 2 或 4

7.4 PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:64

作用：

• 缓解显存碎片问题

规律：

• 有时能减少“明明还有一点显存却申请失败”的情况
• 不能代替降低 batch size

结论：

• 它是辅助措施，不是主修复手段

---

8. 建议的排障顺序

训练报错后，不要盲目改很多地方。推荐按下面顺序排查。

第一步：先看错误类型

如果是：

• CUDA out of memory

那就优先怀疑：

• mini_batch_size
• n_sample_points

不要先去改数据集代码。

第二步：先用命令行覆盖参数

优先尝试：

• --mini_batch_size 2
• --n_sample_points 8192

如果还是 OOM，再继续降。

第二步半：如果降点数后出现 Sizes of tensors must match

优先检查两件事：

1. DenseFusion 是否仍在使用固定 AvgPool1d(num_points)
2. LM_Dataset 是否共享了训练脚本里覆盖后的 config

如果这两点没修，就很容易在降点数后出现新的形状错误。

第三步：确认是不是碎片问题

如果日志里提示：

If reserved memory is >> allocated memory try setting max_split_size_mb

那就加：

PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:64

第四步：再看是否有其他代码警告

例如：

• lr_scheduler.step() 调用顺序问题
• 过时 API 警告

这些不一定致命，但应该尽量修掉，避免训练行为偏离预期。

---

9. 常见误区

9.1 误区一：OOM 一定是模型有 bug

不一定。

很多时候只是默认配置是为更大显存卡准备的。当前仓库默认的：

• batch_size=24
• n_sample_points=12288

对 8 GiB 左右显卡很容易超。

9.2 误区二：只调 num_workers 就能解决 OOM

通常不能。

num_workers 主要影响：

• CPU 负载
• 数据预取速度
• 系统内存占用

而 GPU OOM 的核心仍然是：

• batch 大小
• 点数
• 模型规模

9.3 误区三：只靠 max_split_size_mb 就够了

也不够。

显存碎片最多是次要因素。根本问题还是：

• 当前每步训练需求超过了显卡容量

9.4 误区四：警告可以完全忽略

不能一概忽略。

比如 lr_scheduler.step() 顺序不对，虽然不一定导致训练崩溃，但会导致学习率计划与预期不一致。

9.5 误区五：把 --n_sample_points 传给训练脚本就一定生效

不一定。

如果数据集内部重新创建了自己的 Config，那么你虽然在训练脚本里改了点数，数据集实际仍可能按旧点数输出。

这会造成：

• OOM 问题表面上在降点数
• 实际输入和模型期望点数不一致
• 最终转成 DenseFusion 拼接时报错

---

10. 推荐的最小稳定训练方案

如果你的 GPU 只有大约 8 GiB 显存，建议从下面这组配置开始：

conda activate pvn3d
cd /workspace/workflow/self/PVN3D/pvn3d
PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:64 python -m train.train_linemod_pvn3d \
  --cls ape \
  --mini_batch_size 1 \
  --val_mini_batch_size 1 \
  --n_sample_points 4096 \
  --num_workers 2

如果这组能稳定跑起来，再逐步尝试往上加：

1. 先把 n_sample_points 提到 8192
2. 再把 mini_batch_size 提到 2

不要一次同时把多个参数抬高，否则很难判断到底是谁导致再次 OOM。

如果你已经应用了本仓库里的修复，那么这组命令也同时规避了：

• 小显存 OOM
• DenseFusion 固定池化引起的形状不匹配
• LM_Dataset 与训练脚本配置不一致

---

11. 训练前自检清单

开始训练前，建议确认以下几点：

• 已进入 pvn3d-dev 容器
• 已执行 conda activate pvn3d
• 当前路径是 /workspace/workflow/self/PVN3D/pvn3d
• 数据集路径和 train.txt、test.txt 正常
• GPU 显存容量已确认
• 已根据显存设置合适的 mini_batch_size
• 已根据显存设置合适的 n_sample_points
• 如果之前报过碎片相关 OOM，已加 PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:64

---

12. 总结

这次训练失败的主因是：

• 当前 GPU 显存太小，无法承受 PVN3D 默认训练配置

直接触发 OOM 的关键默认参数是：

• mini_batch_size = 24
• n_sample_points = 12288

正确的处理方式不是盲目改模型，而是先做以下动作：

1. 降低 mini_batch_size
2. 降低 n_sample_points
3. 必要时加入 PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:64
4. 顺手修正 lr_scheduler.step() 的调用顺序
5. 确认模型和数据集使用的是同一份点数配置
6. 将 DenseFusion 的全局池化改为自适应池化

如果后续仍然报错，再根据新的日志区分是：

• 继续 OOM
• 降点数后出现形状不匹配
• 数据问题
• checkpoint 问题
• 其他代码逻辑问题

按这个顺序排查，效率会明显高于一开始就大面积改代码。