PVN3D 模型训练 Bug 调试指南

豆破苍琼

295人浏览 · 2026-03-25 11:37:47

豆破苍琼 · 2026-03-25 11:37:47 发布

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 主干网络的卷积阶段触发的显存不足。

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

lr_scheduler.step() 在 optimizer.step() 之前调用
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 的警告，因为推荐顺序是：

optimizer.step()
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. 当前仓库已做的修复

本次已修改训练脚本：

5.1 新增可覆盖的训练参数

新增了以下命令行参数：

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

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

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

将训练循环调整为：

loss.backward()
optimizer.step()
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` 复用训练脚本中的配置

位置：

修复前的问题是：

训练脚本通过命令行修改了 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`

优先检查两件事：

DenseFusion 是否仍在使用固定 AvgPool1d(num_points)
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