模型转换成移动端格式

有时候，需要将在电脑端训练好的模型集成到移动端应用中，在移动端以离线方式运行。这时候，模型文件的格式会与电脑端需要的格式有些区别。移动端应用有Android原生、IOS原生程序，有flutter应用程序。Android原生程序可以将.pt格式模型进行集成； flutter应用程序可以通过​flutter_pytorch​ 插件集成.pt格式模型； IOS原生程序可以集成.mlmodel格式的模型文件，这就要求将模型由.pt格式转换成.mlmodel格式。

下面分别介绍在Pytorch中如何把模型由.pth转换成.pt与由.pt转换成.mlmodel格式。

模型由.pth转换成.pt格式：

import torch

from resolution_model import SuperResolutionNet

import torch.utils.model_zoo as model_zoo

mobile_pt ='superres_epoch100.pt' # 将模型保存为Android可以调用的文件

torch_model = SuperResolutionNet(upscale_factor=3)

#加载预训练的模型权重

model_url = 'https://s3.amazonaws.com/pytorch/test_data/export/superres_epoch100-44c6958e.pth'

batch_size = 1 # just a random number

# 从预训练的模型权重初始化模型

map_location = lambda storage, loc: storage

if torch.cuda.is_available():

map_location = None

torch_model.load_state_dict(model_zoo.load_url(model_url, map_location=map_location))

torch_model.eval()

input_tensor = torch.rand(batch_size, 1, 224, 224, requires_grad=True) # 设定输入数据格式

torch_model = torch.jit.trace(torch_model, input_tensor) # 模型转化

torch_model.save(mobile_pt) # 保存文件

运行程序后转换成.pt模型，如下图所示：

模型由.pt转换成.mlmodel格式：

先在Python虚拟环境中安装coremltools：

source /opt/venv/python39venv_convert_model/bin/activate
pip install coremltools

模型由.pt转换成.mlmodel格式代码如下 ：

import torch
import platform
import coremltools as ct
from torch.nn import functional as F
from coremltools.converters.mil.mil import types
from resolution_model import SuperResolutionNet  # your model

# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device = torch.device("cpu")

# 加载TorchScript模型
model = torch.jit.load('superres_epoch100.pt')

# 确保模型在评估模式
model.eval()
# 定义输入信息
input_shape = (1, 1, 224, 224)  # 根据你的模型输入形状调整
input_features = [
    ct.TensorType(name="input_1", shape=input_shape)
]
# 转换模型到CoreML
mlmodel = ct.convert(
    model,
    inputs=input_features,
    minimum_deployment_target=ct.target.iOS14  # 根据需要调整
)
# 保存CoreML模型
mlmodel.save('superres_epoch100.mlmodel')
运行程序后生成.mlmodel格式的模型文件，如下图所示：