如何在LibTorch中实现模型间的参数复制？

你在PyTorch中用net2.load_state_dict(net.state_dict())实现参数复制的思路完全正确，LibTorch也提供了几乎一致的便捷API，同时也支持手动实现参数复制。先分析你现有代码的问题，再给出正确实现：

你的代码存在的问题 #

1. 参数复制方向搞反：你的loadstatedict函数是把target_model的参数复制到model中，但实际需求是将model（源）的参数复制到target_model（目标），和PyTorch的逻辑相反。
2. 未递归获取目标模型参数：target_model.named_parameters()默认不递归子模块，需要传入true才能拿到所有层级的参数。
3. 模块注册不规范：自定义模型中的lin1、lin2最好显式调用register_module注册，避免LibTorch无法正确追踪参数。

方法一：使用LibTorch官方load_state_dict（最便捷） #

和PyTorch用法几乎完全一致，源模型调用state_dict()获取参数字典，目标模型调用load_state_dict()加载：

#include <torch/torch.h>

torch::Device device(torch::kCUDA);

struct Critic_Net : torch::nn::Module {
public:
    Critic_Net() {
        // 显式注册子模块，确保参数能被正确追踪
        lin1 = register_module("lin1", torch::nn::Linear(3, 3));
        lin2 = register_module("lin2", torch::nn::Linear(3, 1));
        lin1->to(device);
        lin2->to(device);
    }

    torch::Tensor forward(torch::Tensor x) {
        auto h = torch::relu(lin1->forward(x));
        return lin2->forward(h);
    }

private:
    torch::nn::Linear lin1, lin2;
};

int main() {
    auto net = Critic_Net();
    auto net2 = Critic_Net();
    auto the_ones = torch::ones({3, 3}).to(device);

    std::cout << "复制前net输出:\n" << net.forward(the_ones) << "\n";
    std::cout << "复制前net2输出:\n" << net2.forward(the_ones) << "\n";

    // 核心代码：复制参数
    net2.load_state_dict(net.state_dict());

    std::cout << "复制后net输出:\n" << net.forward(the_ones) << "\n";
    std::cout << "复制后net2输出:\n" << net2.forward(the_ones) << "\n";

    return 0;
}

方法二：手动实现参数复制（修正你的函数） #

如果需要自定义复制逻辑，调整复制方向并确保递归获取所有参数和buffer：

#include <torch/torch.h>

torch::Device device(torch::kCUDA);

// 将source_model的参数复制到target_model
void copy_model_params(torch::nn::Module& source_model, torch::nn::Module& target_model) {
    torch::autograd::GradMode::set_enabled(false); // 禁用梯度计算，提升复制效率

    // 递归获取源模型的所有参数和buffer
    auto source_params = source_model.named_parameters(true);
    auto source_buffers = source_model.named_buffers(true);

    // 递归获取目标模型的所有参数和buffer
    auto target_params = target_model.named_parameters(true);
    auto target_buffers = target_model.named_buffers(true);

    // 复制参数
    for (const auto& [name, src_tensor] : source_params) {
        if (auto* tgt_tensor = target_params.find(name)) {
            tgt_tensor->copy_(src_tensor);
        }
    }

    // 复制buffer（如BN的running_mean等）
    for (const auto& [name, src_tensor] : source_buffers) {
        if (auto* tgt_tensor = target_buffers.find(name)) {
            tgt_tensor->copy_(src_tensor);
        }
    }
}

struct Critic_Net : torch::nn::Module {
public:
    Critic_Net() {
        lin1 = register_module("lin1", torch::nn::Linear(3, 3));
        lin2 = register_module("lin2", torch::nn::Linear(3, 1));
        lin1->to(device);
        lin2->to(device);
    }

    torch::Tensor forward(torch::Tensor x) {
        auto h = torch::relu(lin1->forward(x));
        return lin2->forward(h);
    }

private:
    torch::nn::Linear lin1, lin2;
};

int main() {
    auto net = Critic_Net();
    auto net2 = Critic_Net();
    auto the_ones = torch::ones({3, 3}).to(device);

    std::cout << "复制前net输出:\n" << net.forward(the_ones) << "\n";
    std::cout << "复制前net2输出:\n" << net2.forward(the_ones) << "\n";

    // 调用自定义复制函数
    copy_model_params(net, net2);

    std::cout << "复制后net输出:\n" << net.forward(the_ones) << "\n";
    std::cout << "复制后net2输出:\n" << net2.forward(the_ones) << "\n";

    return 0;
}

关键注意事项 #

• 确保源模型和目标模型结构完全一致（包括子模块名称、参数形状），否则会出现参数匹配失败的问题。
• 若模型在GPU上运行，需确保源和目标模型都已移动到同一设备，避免复制时设备不匹配的错误。
• 显式注册子模块（register_module）是LibTorch的规范做法，能确保模型的参数、buffer被正确追踪，避免state_dict遗漏参数。

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