olivame
e1b029201e
目前参数--encoder_mode有三种选择: 1. fp32: 全部使用fp32, 适合显存充足的情况 2. autocast: 使用torch.cuda.amp.autocast自动混合精度, 稍微快一些, psnr下降较多 3. bf16_full: 全部使用bf16, 精度较高
109 lines
6.3 KiB
Bash
109 lines
6.3 KiB
Bash
python3 psnr_score_for_challenge.py --gt_video unitree_g1_pack_camera/case1/unitree_g1_pack_camera_case1.mp4 --pred_video unitree_g1_pack_camera/case1/output/inference/0_full_fs6.mp4 --output_file unitree_g1_pack_camera/case1/psnr_result.json
|
||
采样阶段（synthesis/ddim_sampling）+ UNet 前向几乎吃掉全部时间，外加明显的 CPU 侧 aten::to/_to_copy 和 aten::copy_ 开销；整体优化优先级还是“减少采样步数 + 加速每步前向 + 降低无谓拷贝”。下
|
||
面是更针对这份 profile 的思路：
|
||
|
||
- 优先级1：减少采样步数/换更快采样器
|
||
- 把 DDIM 30 步降到 10–20 步，或改用 DPM-Solver++/UniPC；这往往是 1.5–3× 的最直接收益。采样逻辑在 src/unifolm_wma/models/samplers/ddim.py，入口在 scripts/evaluation/
|
||
world_model_interaction.py。
|
||
- 如允许训练侧投入，可做蒸馏（LCM/Consistency）让 4–8 步也可用。
|
||
- 优先级1：每步前向加速（编译 + AMP + TF32）
|
||
- torch.compile 只包 diffusion_model，见 scripts/evaluation/world_model_interaction.py（你文档里也已写）。
|
||
- 推理包一层 autocast（fp16/bf16）并开启 TF32：
|
||
torch.backends.cuda.matmul.allow_tf32 = True + torch.set_float32_matmul_precision("high")。
|
||
- torch.backends.cudnn.benchmark = True 在固定 shape 下很有效。
|
||
- 优先级1：消除循环内的 to()/copy/clone
|
||
- profile 里 aten::_to_copy CPU 时间很高，建议逐步排查是否在 DDIM loop 或条件准备里重复 .to(device) / .float() / .half()。
|
||
- 把常量（timesteps/sigmas/ts 等）提前放 GPU，避免每步创建；避免不必要的 clone()。
|
||
- 重点排查 src/unifolm_wma/models/samplers/ddim.py 与 scripts/evaluation/world_model_interaction.py 的数据准备段。
|
||
- 优先级2：注意力实现确认
|
||
- attention 只占 3% 左右，但如果没启用 xformers/SDPA，仍有收益空间。检查 src/unifolm_wma/modules/attention.py 的 XFORMERS_IS_AVAILBLE。
|
||
- 无 xformers 时可改用 scaled_dot_product_attention（Flash Attention 路径）。
|
||
- 优先级2：VAE 解码 & 保存 I/O
|
||
- synthesis/decode_first_stage 仍是秒级，建议 autocast + 可能的 torch.compile。位置在 src/unifolm_wma/models/autoencoder.py。
|
||
- save_results 约 38s：如果只是评测，考虑降低保存频率/分辨率或异步写盘。
|
||
- 优先级3：结构性减负
|
||
- 降低 temporal_length、输入分辨率或 model_channels 会线性降低 compute（配置在 configs/inference/world_model_interaction.yaml）。
|
||
- 如果 action_generation 与 world_model_interaction 共享条件，可以缓存 CLIP/VAE 编码，避免重复计算（model_architecture_analysis.md 的条件编码流程已说明）。
|
||
|
||
如果你希望我直接落地改动，推荐顺序：
|
||
|
||
1. torch.compile + AMP + TF32 + cudnn.benchmark
|
||
2. 排查 .to()/copy/clone 的重复位置并移出循环
|
||
3. 若需要更大幅度，再换采样器/降步数
|
||
|
||
|
||
|
||
|
||
|
||
A100 上我推荐 BF16 优先（稳定性更好、PSNR 更稳），FP16 作为速度优先方案。
|
||
|
||
下面是“分模块”的 消融方案（从稳到激进）：
|
||
|
||
0）基线
|
||
|
||
- 全 FP32（你现在就是这个）
|
||
|
||
1）只对扩散主干做 BF16（最推荐）
|
||
|
||
- 量化对象：model.model（扩散 UNet/WMAModel 主体）
|
||
- 保持 FP32：first_stage_model（VAE 编/解码）、cond_stage_model（文本）、embedder（图像）、image_proj_model
|
||
- 预期：PSNR 基本不掉 or 极小波动
|
||
|
||
2）+ 轻量投影/MLP 做 BF16
|
||
|
||
- 增加：image_proj_model、state_projector、action_projector
|
||
- 预期：几乎不影响 PSNR
|
||
|
||
3）+ 文本/图像编码做 BF16
|
||
|
||
- 增加：cond_stage_model、embedder
|
||
- 预期：可能有轻微波动，通常仍可接受
|
||
|
||
4）VAE 也做 BF16（最容易伤 PSNR）
|
||
|
||
- 增加：first_stage_model
|
||
- 预期：画质/PSNR 最敏感，建议最后做消融
|
||
|
||
———
|
||
|
||
具体建议（A100）
|
||
|
||
- 优先 BF16：稳定性好于 FP16
|
||
- 只做半精度，不做 INT 量化：保持 PSNR
|
||
- VAE 尽量 FP32：最影响画质的模块
|
||
|
||
BF16 projector比FP32 projector更准的可能原因：
|
||
- 数值路径更一致：主干在 BF16 下做 attention/MLP，projector 若是 FP32 会在进入主干前被 downcast，导致“先高精度非线性→再截断”的分布偏移；直接 BF16 算 projector 反而让输出
|
||
分布更贴近主干的计算习惯。
|
||
- 训练分布匹配：训练时你是 precision:16，projector 长期在低精度环境下被优化；推理用 FP32 反而可能偏离训练时的统计特性。
|
||
- LayerNorm/Softmax 敏感：Resampler/MLP 里 LN/Softmax 对精度很敏感，FP32 计算后再降精度，数值边界更容易“硬截断”；BF16 全程计算可能更平滑。
|
||
|
||
这也解释了为什么你看到 BF16 projector 反而更准。
|
||
|
||
embedder：
|
||
改成 autocast only（权重 FP32，预处理 FP32，仅主干 BF16）
|
||
|
||
- 效果差的那次：encoder autocast set to bf16 (weights=fp32)
|
||
也就是“权重 FP32 + autocast BF16（我现在的 autocast only 方案）”
|
||
- 效果好的那次：cond_stage_model / embedder param dtype: torch.bfloat16
|
||
也就是“权重 BF16（全 BF16）”
|
||
|
||
为什么 FP32 反而掉分大？主要是“混精度路径不一致”：
|
||
|
||
1. autocast only 不是等价于 full BF16
|
||
autocast 会把部分算子留在 FP32（比如 LayerNorm / Softmax），另一部分用 BF16；
|
||
这会让 encoder 输出分布和“全 BF16”不一样，尤其当扩散主干也是 BF16 时，条件向量分布一偏就会影响 PSNR。
|
||
2. 先 FP32 计算再转 BF16
|
||
encoder 的输出最后还是要拼到 BF16 主干里，等于“先高精度 → 再截断”。
|
||
训练时如果你用的是 BF16，模型更适应“低精度直接计算”的分布。
|
||
3. 你的现象和 projector 一样
|
||
你已经观察到 projector BF16 更准，这说明模型很可能“更偏好一致的 BF16 路径”。
|
||
|
||
结论：
|
||
对你这个模型，full BF16（权重 BF16）更匹配训练分布，反而 PSNR 更高。
|
||
|
||
如果你要我改回“bf16=全 BF16（权重 BF16）”，并且保留现在的 autocast only 作为可选模式，我可以这样做：
|
||
|
||
1. 新增 --encoder_mode {fp32, autocast, bf16_full}
|
||
2. bf16_full = 权重 BF16 + 前向 BF16
|
||
3. autocast = 权重 FP32 + 仅主干 autocast（现在的实现） |