【Libtorch部署】pkl 转 pt

最新推荐文章于 2025-04-27 19:46:22 发布

Xnion

最新推荐文章于 2025-04-27 19:46:22 发布

阅读量1.2k

点赞数

CC 4.0 BY-SA版权

文章标签： pytorch c++ opencv

本文链接：https://blog.youkuaiyun.com/Xnion/article/details/126929387

本文介绍了如何将PyTorch的pkl格式模型文件转换为pt格式，以便于在C++环境中部署Libtorch使用。提供了两种不同的转换方法，并强调了在不同环境下（CPU/GPU）模型的具体操作细节。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

在C++部署Libtorch的时候，需要调用 pt 格式的模型，因此需要提前将 pkl 格式的模型转成 pt 格式。

CPU 版本

import torch
from myModels.UNet import UNet as unet

# Download and load the pre-trained model
# 模型加载方式1
# model = torch.load("checkpoints/model_50.pkl",map_location='cpu')

# 模型加载方式2，需要提前加载模型框架（断点）
num_channels = 3
num_classes = 1
model = unet(n_classes=num_classes,n_channels=num_channels).to(device='cpu')

checkpoint = torch.load('./checkpoints(Unet)/model.pkl')            # 加载断点
model.load_state_dict(checkpoint['net'])            # 加载模型可学习参数

# 模型加载方式3，需要提前加载模型框架（非断点）
# model.load_state_dict(torch.load('checkpoints/model_50.pkl'))

model.eval()

example_input = torch.rand(1, 3, 256, 256)
script_module = torch.jit.trace(model, example_input)
script_module.save('./checkpoints(Unet)/model_cpu.pt')

- GPU 版本（看看就好，其实用不着）

import torch
from myModels.UNet import UNet as unet
# 模型加载
num_channels = 3
num_classes = 1
model = unet(n_classes=num_classes,n_channels=num_channels).to(device='cuda')

checkpoint = torch.load('./checkpoints(Unet)/model.pkl')            # 加载断点
model.load_state_dict(checkpoint['net'])            # 加载模型可学习参数

model.cuda()
model.eval()

example_input = torch.rand(1, 3, 256, 256).cuda()  # 生成测试sample
script_module = torch.jit.trace(model, example_input)
script_module.save('./checkpoints(Unet)/model_gpu.pt')

【P.S.】

Python中无需将模型转成GPU版本的 pt，因为CPU版本的 pt 在C++中Release模型下也可以调用CUDA，相反GPU版本的 pt 在C++中Debug模型还会报错。
在Libtorch中：Debug 模型下，默认设备使用的是 CPU； Releas 模型下，根据用户设置使用 GPU / CPU。