【Libtorch部署】pkl 转 pt

在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。