Libtorch(C++)：保存/加载模型与优化器状态字典以重训的高效正确方法

核心问题 #

我需要确认在LibTorch中同时保存、加载模型与优化器状态字典并进行重训的正确流程，尤其是创建Optimizer并传入模型参数的时机——应该在加载状态字典之前、之后执行，还是两者无关？

现有代码示例 #

保存模型与优化器状态 #

// Save model state
torch::serialize::OutputArchive output_model_archive;
myModel.to(torch::kCPU);
myModel.save(output_model_archive);
output_model_archive.save_to(model_state_dict_path);
// Save optim state
torch::serialize::OutputArchive output_optim_archive;
myOptimizer->save(output_optim_archive);
output_optim_archive.save_to(optim_state_dict_path);

重训时加载状态 #

// Load model state
torch::serialize::InputArchive input_archive;
input_archive.load_from(state_dict);
myModel.load(input_archive);
// Load optim state
torch::serialize::InputArchive input_archive; // 注意：这里变量名重复，会编译报错
input_archive.load_from(state_dict);
myOptimizer->load(input_archive);

创建Optimizer对象 #

std::shared_ptr<torch::optim::Optimizer> myOptimizer;
myOptimizer.reset(new torch::optim::Adam(myModel.parameters(), torch::optim::AdamOptions(LR)));

正确流程与关键说明 #

优化器创建时机：必须在加载状态字典之前 #

优化器的状态字典（如Adam的动量缓存exp_avg、exp_avg_sq）是与模型参数一一绑定的，只有当优化器已经持有模型参数的引用时，加载状态字典才能正确将缓存关联到对应的参数上。如果先加载模型状态再创建优化器，优化器会初始化全新的状态，直接覆盖你加载的内容；反之，先创建优化器再加载模型状态，模型参数的更新会自动同步到优化器持有的引用中，后续加载优化器状态就能正确恢复之前的训练状态。

修正后的完整重训流程 #

// 1. 初始化模型实例
MyModel myModel;

// 2. 创建优化器（关键步骤：必须在加载状态前绑定模型参数）
std::shared_ptr<torch::optim::Optimizer> myOptimizer = 
    std::make_shared<torch::optim::Adam>(myModel.parameters(), torch::optim::AdamOptions(LR));

// 3. 加载模型状态
torch::serialize::InputArchive model_archive;
model_archive.load_from(model_state_dict_path);
myModel.load(model_archive);

// 4. 加载优化器状态（注意使用不同的变量名，避免编译错误）
torch::serialize::InputArchive optim_archive;
optim_archive.load_from(optim_state_dict_path);
myOptimizer->load(optim_archive);

// 5. （可选）将模型与优化器转移到目标训练设备（如CUDA）
myModel.to(torch::kCUDA);
// 同步优化器状态到目标设备
for (auto& group : myOptimizer->param_groups()) {
    for (auto& p : group.params()) {
        p.data() = p.data().to(torch::kCUDA);
        if (p.grad().defined()) {
            p.grad().data() = p.grad().data().to(torch::kCUDA);
        }
    }
    // 针对Adam等带额外状态的优化器，同步缓存到目标设备
    for (auto& p : group.params()) {
        auto& state = myOptimizer->state()[p];
        if (state.contains("exp_avg")) {
            state["exp_avg"] = state["exp_avg"].to(torch::kCUDA);
        }
        if (state.contains("exp_avg_sq")) {
            state["exp_avg_sq"] = state["exp_avg_sq"].to(torch::kCUDA);
        }
    }
}

// 6. 开始重训
for (int epoch = 0; epoch < num_epochs; ++epoch) {
    // 前向传播
    auto output = myModel(input.to(torch::kCUDA));
    auto loss = criterion(output, target.to(torch::kCUDA));
    
    // 反向传播与优化步骤
    myOptimizer->zero_grad();
    loss.backward();
    myOptimizer->step();
}

// 7. 保存状态（与原逻辑一致，先移模型到CPU避免设备不兼容）
myModel.to(torch::kCPU);
torch::serialize::OutputArchive model_output;
myModel.save(model_output);
model_output.save_to(model_state_dict_path);

torch::serialize::OutputArchive optim_output;
myOptimizer->save(optim_output);
optim_output.save_to(optim_state_dict_path);

关键注意事项 #

• 避免变量名重复：加载模型和优化器状态时，不能重复定义InputArchive变量，否则会触发C++编译错误。
• 设备匹配：保存前将模型移到CPU是最佳实践，避免不同设备（CPU/CUDA）之间的状态加载冲突；加载后需要手动将模型和优化器状态同步到目标训练设备。
• 优化器与模型参数绑定：优化器必须绑定当前训练的模型参数，不能复用其他模型的优化器状态，否则会导致参数不匹配。

内容的提问来源于stack exchange，提问作者Brandon