pytorch 如何在LibTorch中将一个模型的参数复制到另一个模型?

6ie5vjzr  于 2022-11-09  发布在  其他
关注(0)|答案(1)|浏览(281)

如何在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);

}
yhxst69z

yhxst69z1#

如果我没理解错的话,你的load_state_dict解决方案应该可以用。这里的问题和你之前的问题是一样的:没有注册任何参数、缓冲区或子模块。添加register_module调用,它应该可以正常工作。
Link to the question
此类别的外观:

struct Critic_Net : torch::nn::Module {
public:
    Critic_Net() {
        lin1 = register_module("lin1", torch::nn::Linear(427, 42));
        lin2 = register_module("lin1", torch::nn::Linear(42, 286));
        lin3 = register_module("lin1", torch::nn::Linear(286, 1));
    }
     torch::Tensor forward(torch::Tensor next_state_batch__sampled_action) {
     // unchanged 
    }
    torch::nn::Linear lin1{nullptr}, lin2{nullptr}, lin3{nullptr};
};

相关问题