pytorch从python转 c++涉及到的数据保存加载问题；libtorch

最新推荐文章于 2025-06-12 13:57:47 发布

原创最新推荐文章于 2025-06-12 13:57:47 发布 · 732 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#python #pytorch #c++

外功-算法与计算专栏收录该内容

5 篇文章

订阅专栏

博客聚焦于PyTorch从Python转C++时的数据保存加载问题。一是torch.nn.Module保存的state_dict无法被C++访问，需转化为Python字典；二是介绍了保存Tensor张量及C++加载Tensor张量的相关内容，并给出了Python和C++代码示例。

部署运行你感兴趣的模型镜像

pytorch 从python 转 c++涉及到的数据保存加载问题

1. torch.nn.Module 保存 state_dict 无法被 c++访问，只能转化为 python字典

python代码

model = ThreeLayer_FCNN_Net()
model.load_state_dict(ret_load)
w = {k: v for k, v in model.state_dict().items()}
torch.save(w, "data.pkl")

c++代码

std::vector<char> get_the_bytes(std::string filename) {
    std::ifstream input(filename, std::ios::binary);
    std::vector<char> bytes(
        (std::istreambuf_iterator<char>(input)),
        (std::istreambuf_iterator<char>()));
    input.close();
    return bytes;
}

std::vector<char> f = get_the_bytes("../48_channal_guozi_2.pkl");
torch::IValue x = torch::pickle_load(f);
auto dict = x.toGenericDict();
std::cout << "dict: " << dict << std::endl;

2. 保存 Tensor张量，C++加载 Tensor张量

python代码

data = torch.ones(3, 4)
torch.save(data, "data.pkl")

c++代码

std::vector<char> f = get_the_bytes("../48_channal_guozi_2.pkl");
torch::IValue x = torch::pickle_load(f);
auto my_tensor = x.toTensor();
std::cout << "my_tensor: " << my_tensor << std::endl;