Open-Sora-Plan项目中微调模型推理时的模型选择指南

在Open-Sora-Plan项目中，当用户完成模型微调后，常常会遇到如何选择正确的模型文件进行推理的问题。本文将详细介绍微调后生成的模型文件结构，以及如何正确选择用于推理的模型文件。

微调后的模型文件结构

典型的微调完成后，会生成以下目录结构：

checkpoint
├── model_ema
│   ├── diffusion_pytorch_..._fetensors.index.json
│   ├── diffusion_pytorch_..._f-00002.safetensors
│   ├── diffusion_pytorch_..._f-00002.safetensors
│   └── config.json
├── pytorch_model
│   ├── mp_rank_00_model_states.pt
│   ├── bf16_zero_pp_rank_..._00_optim_states.pt
│   └── bf16_zero_pp_rank_..._00_optim_states.pt
└── model
    ├── config.json
    ├── diffusion_pytorch_model.safetensors
    ├── zero_to_fp32.py
    ├── scheduler.bin
    ├── random_states_1.pkl
    ├── random_states_0.pkl
    └── latest

各目录作用解析

1. model_ema目录：包含指数移动平均(EMA)版本的模型权重，通常能提供更稳定的推理结果
2. pytorch_model目录：包含分布式训练时的模型状态和优化器状态，主要用于训练恢复
3. model目录：包含标准版本的微调后模型，以及相关的配置文件

推理时的模型选择建议

根据项目官方建议，进行推理时可以从以下两种模型中选择：

1. model_ema目录下的模型：

   • 优势：EMA模型通常能提供更稳定的输出
   • 文件选择：使用目录中的safetensors文件和配套的config.json

2. model目录下的模型：

   • 优势：这是标准训练过程的直接产物
   • 文件选择：使用diffusion_pytorch_model.safetensors和config.json

技术细节说明

• EMA模型：在训练过程中维护模型权重的移动平均，有助于平滑训练过程中的波动，通常能获得更好的泛化性能
• 标准模型：直接反映最后一次训练迭代的模型状态
• VAE选择：在Open-Sora-Plan项目中，VAE通常不需要单独选择，因为模型已经包含了完整的扩散模型结构

实践建议

对于大多数应用场景，建议优先尝试使用model_ema目录下的模型进行推理，因为它通常能提供更稳定的结果。如果遇到特定需求或问题，可以再尝试使用model目录下的标准模型进行对比。

通过理解这些模型文件的作用和选择原则，用户可以更有效地利用微调后的模型进行推理任务。