pytorch模型在IOS下部署

由于实际情况考虑,这里贴出官方那个代码:

import torch
import torchvision

# Load a pre-trained version of MobileNetV2
torch_model = torchvision.models.mobilenet_v2(pretrained=True)
# Set the model in evaluation mode
torch_model.eval()
# Trace with random data
example_input = torch.rand(1, 3, 224, 224) # after test, will get size mismatch error message with size 256x256
traced_model = torch.jit.trace(torch_model, example_input)
# Download class labels (from a separate file)
import urllib
label_url = https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt
class_labels = urllib.request.urlopen(label_url).read().decode("utf-8").splitlines()

class_labels = class_labels[1:] # remove the first class which is background
assert len(class_labels) == 1000
import coremltools as ct
# Convert to Core ML using the Unified Conversion API
model = ct.convert(
    traced_model,
    inputs=[ct.ImageType(name="input_1", shape=example_input.shape)], #name "input_1" is used in quickstart
    classifier_config = ct.ClassifierConfig(class_labels) # provide only if step 2 was performed
)
# Save model
model.save("MobileNetV2.mlmodel")

预测的话需要在Mac设备上进行,只需要调取模型,然后做一个prediction。 转换之后的Core ML模型可以记载进入xcode之下,然后进行输入输出的确定及就可以在移动设备上运行。

经验分享 程序员 微信小程序 职场和发展