coremltools_pytorch_mobilenetv2_conversion_and_prediction.py

1import torch
2import torchvision
3import coremltools as ct
4
5# 1. Load a pre-trained version of MobileNetV2 from torchvision
6torch_model = torchvision.models.mobilenet_v2(pretrained=True)
7torch_model.eval()
8
9# 2. Trace the model with dummy input to produce a TorchScript object
10example_input = torch.rand(1, 3, 224, 224)
11traced_model = torch.jit.trace(torch_model, example_input)
12
13# 3. Convert the model to Core ML using the coremltools converter
14model = ct.convert(
15    traced_model,
16    inputs=[ct.TensorType(shape=example_input.shape)]
17)
18
19# 4. Save the converted model
20model.save("MobileNetV2.mlpackage")
21
22# 5. (Optional) Run a prediction on the Core ML model
23# Note: Prediction on Mac requires macOS and is performed via the Core ML framework
24prediction = model.predict({"input_1": example_input.numpy()})
25print(prediction)