如何在 LibTorch 中将一个模型的参数复制到另一个模型？答案

【问题标题】：How can I copy the parameters of one model to another in LibTorch?如何在 LibTorch 中将一个模型的参数复制到另一个模型？
【发布时间】：2022-10-25 20:35:02
【问题描述】：

如何在 LibTorch 中将一个模型的参数复制到另一个模型？我知道如何在 Torch (Python) 中做到这一点。

net2.load_state_dict(net.state_dict())

我已经在 C++ 中尝试了下面的代码，并做了很多工作。它没有复制一个到另一个。
我没有看到将一个类似模型的参数复制到另一个类似模型的选项。

#include <torch/torch.h>
using namespace torch::indexing;

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

void loadstatedict(torch::nn::Module& model, torch::nn::Module& target_model) {
torch::autograd::GradMode::set_enabled(false);  // make parameters copying possible
auto new_params = target_model.named_parameters(); // implement this
auto params = model.named_parameters(true /*recurse*/);
auto buffers = model.named_buffers(true /*recurse*/);
for (auto& val : new_params) {
    auto name = val.key();
    auto* t = params.find(name);
    if (t != nullptr) {
        t->copy_(val.value());
    } else {
        t = buffers.find(name);
        if (t != nullptr) {
            t->copy_(val.value());
            }
        }
    }
}

struct Critic_Net : torch::nn::Module {
    torch::Tensor next_state_batch__sampled_action;
    public:
    Critic_Net() {
        lin1 = torch::nn::Linear(3, 3);
        lin2 = torch::nn::Linear(3, 1);
        lin1->to(device);
        lin2->to(device);
    }
    torch::Tensor forward(torch::Tensor next_state_batch__sampled_action) {
        auto h = next_state_batch__sampled_action;
        h = torch::relu(lin1->forward(h));
        h = lin2->forward(h);
        return h;
    }
    torch::nn::Linear lin1{nullptr}, lin2{nullptr};
};

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

int main() {
    std::cout << net.forward(the_ones);
    std::cout << net2.forward(the_ones);
    loadstatedict(net, net2);
    std::cout << net.forward(the_ones);
    std::cout << net2.forward(the_ones);

}

【问题讨论】：

标签： neural-network pytorch libtorch

【解决方案1】：

如果我理解正确，您使用load_state_dict 的解决方案应该可以工作。这里的问题与您之前的问题相同：没有任何内容被注册为参数、缓冲区或子模块。添加register_module 调用，它应该可以正常工作。

Link to the question

【讨论】：