解决metahuman-stream项目中模型加载时的shape mismatch问题

在使用metahuman-stream项目进行数字人训练时，用户遇到了一个常见的模型加载错误：size mismatch。这个错误通常发生在尝试加载预训练模型或不同训练配置生成的模型时。本文将详细分析这个问题的原因，并提供完整的解决方案。

问题现象

用户在使用metahuman-stream训练自己的模型后，尝试加载时遇到了以下错误：

RuntimeError: Error(s) in loading state_dict for NeRFNetwork:
        size mismatch for individual_codes: copying a param with shape torch.Size([12000, 4]) from checkpoint, the shape in current model is torch.Size([10000, 4]).
        size mismatch for individual_codes_torso: copying a param with shape torch.Size([12000, 8]) from checkpoint, the shape in current model is torch.Size([10000, 8]).

这表明模型期望的tensor形状与实际加载的模型参数形状不匹配，具体表现在individual_codes和individual_codes_torso这两个参数上。

问题原因分析

1. 训练配置不一致：metahuman-stream项目中，individual_codes的大小与训练时的配置参数有关。默认情况下，项目可能使用了12000的大小，而用户在训练时可能使用了不同的配置，导致生成了10000大小的参数。

2. 数据类型差异：用户可能使用了float32进行训练，而项目默认使用fp16（float16）模式运行。这种数据类型的不匹配也会导致模型加载问题。

3. ASR模型选择：用户使用了简写的--asr_model esperanto参数，而项目作者使用的是完整的cpierse/wav2vec2-large-xlsr-53-esperanto。虽然理论上应该指向同一个模型，但在某些情况下可能导致不一致的行为。

解决方案

方法一：调整训练参数

按照项目作者推荐的完整训练流程进行训练：

# 第一阶段训练
python main.py data/mine-new/ --workspace workspace/trial_mine/ -O --iters 100000 --asr_model cpierse/wav2vec2-large-xlsr-53-esperanto

# 第二阶段微调嘴唇
python main.py data/mine-new/ --workspace trial_mine/ -O --iters 125000 --finetune_lips --patch_size 32 --asr_model cpierse/wav2vec2-large-xlsr-53-esperanto

# 第三阶段训练torso
python main.py data/mine-new/ --workspace trial_mine_torso/ -O --torso --head_ckpt trial_mine/checkpoints/ngp_ep0034.pth --iters 200000 --asr_model cpierse/wav2vec2-large-xlsr-53-esperanto

方法二：调整运行配置

如果已经训练完成，可以尝试以下调整：

1. 关闭fp16模式：在app.py中找到并注释掉opt.fp16 = True这一行，强制使用float32模式运行。

2. 检查模型一致性：确保训练时使用的参数与运行时配置一致，特别是individual_codes相关的参数。

3. 模型大小验证：训练完成后，检查生成的模型文件大小是否与预期一致（约38MB）。如果差异较大，可能需要重新训练。

最佳实践建议

1. 使用完整ASR模型路径：始终使用完整的ASR模型路径，如cpierse/wav2vec2-large-xlsr-53-esperanto，而不是简写形式。

2. 保持训练配置一致：在训练和推理阶段使用相同的配置参数，特别是影响模型结构的参数。

3. 分阶段训练：按照项目作者推荐的三阶段训练流程进行训练，确保每个阶段都正确完成。

4. 验证模型输出：训练完成后，先用main.py进行测试验证，确保模型能正常工作后再集成到应用流程中。

通过以上方法，可以有效解决metahuman-stream项目中模型加载时的shape mismatch问题，确保数字人训练和推理流程的顺利进行。