PyTorch模型转 apple CoreML模型

最新推荐文章于 2024-08-16 09:34:58 发布

JoeyChen1219

最新推荐文章于 2024-08-16 09:34:58 发布

阅读量518

点赞数 1

文章标签： pytorch 深度学习人工智能

本文链接：https://blog.csdn.net/JoeyChen1219/article/details/138131482

版权

本文介绍了如何在苹果Mac上将Torch模型转换为CoreML模型，官方不再推荐ONNX路径，而是推荐直接使用`coremltools`和TorchScript进行模型转换。详细步骤包括安装`coremltools`、使用TorchScript进行模型追踪以及设置输入类型等。

摘要由CSDN通过智能技术生成

如果要部署在苹果mac端，官方已经不推荐使用onnx模型转换CoreML模型了，现在最新支持torch模型直接通过cmtools转换到CoreML模型。

需要安装苹果的CoreML官方工具

pip install -U coremltools

模型转换

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

# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device = torch.device("cpu")
ckpt = torch.load('your_model_path.pt', map_location=device)
model = Model()
model.load_state_dict(ckpt['state_dict'])
model.eval()

Input_1 = torch.randn(1, 3, 256, 256)
Input_2 = torch.randn(1, 3, 256, 256)

ts = torch.jit.trace(model, (Input_1, Input_2))  # TorchScript model

# ct.ImageType
Input_1 = ct.TensorType(name='Input_1', shape=(1, 3, ct.RangeDim(lower_bound=256, upper_bound=2160, default=1280), ct.RangeDim(lower_bound=256, upper_bound=2160, default=736)), dtype=ct.converters.mil.mil.types.fp32)
Input_2 = ct.TensorType(name='Input_2', shape=(1, 3, ct.RangeDim(lower_bound=256, upper_bound=2160, default=1280), ct.RangeDim(lower_bound=256, upper_bound=2160, default=736)), dtype=ct.converters.mil.mil.types.fp32)

ct_model = ct.convert(ts, inputs=[Input_1, Input_2], minimum_deployment_target=ct.target.macOS12)

ct_model.save(your_model_name.mlmodel')