如何在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);
}
1条答案
按热度按时间yhxst69z1#
如果我没理解错的话,你的
load_state_dict
解决方案应该可以用。这里的问题和你之前的问题是一样的:没有注册任何参数、缓冲区或子模块。添加register_module
调用,它应该可以正常工作。Link to the question
此类别的外观: