Convert ClothingCNN with coreai-torch

curl -LsSf https://astral.sh/uv/install.sh | sh
mkdir -p clothing-coreai && cd clothing-coreai
uv init --python 3.11
uv venv
uv add torch torchvision coreai-torch
uv run python -c "import torch, torchvision; print(torch.__version__, torchvision.__version__); print('MPS:', torch.backends.mps.is_available())"
uv run python convert.py

coreai-torch works from a PyTorch exported program. The order matters: load the trained ClothingCNN, export it with an NCHW float32 example, run the Core AI decompositions, then convert. The decomposition table turns composite operations (several steps packed together) into simpler ones the converter can handle.

The conversion script

import torch
from coreai_torch import TorchConverter, get_decomp_table

# Keep the class definition from lesson 02 in scope.
model = ClothingCNN().eval()
model.load_state_dict(torch.load("clothing.pt", map_location="cpu", weights_only=True))

# Match training: float32 NCHW, batch 1, one channel, 28 by 28.
example = torch.zeros(1, 1, 28, 28, dtype=torch.float32)
ep = torch.export.export(model, args=(example,))

# Do this before conversion so composite ops become supported Core AI ops.
ep = ep.run_decompositions(get_decomp_table())

coreai_program = TorchConverter().add_exported_program(
    ep,
    input_names=["image"],
    output_names=["logits"],
).to_coreai()

coreai_program.optimize()
asset = coreai_program.save_asset("clothing.aimodel")
print(asset)

The input and output names become part of the app interface. The Swift function takes image, shaped [1, 1, 28, 28] as float32, and returns logits, shaped [1, 10].

Keep the same image layout

Training and export use NCHW float32. The Swift lesson will build that same layout, so the app does not need an extra NHWC conversion. If a different model or toolchain expects NHWC, follow the descriptor in Xcode's model viewer. Do not guess.

Check the asset

Keep clothing.aimodel after the script finishes. In Xcode's model viewer, check its function, tensor names, shapes, and element types. If your installed coreai-torch release has a Python-side smoke test (a tiny check that it runs), send one zero image through it before you open Xcode. Helper names can change by package version, so use the docs for the package you installed.

Build requirement

When an Xcode target builds with an .aimodel, install the Metal Toolchain in Xcode. Python conversion can succeed while the app build fails if that toolchain is missing.

Sources: coreai-torch documentation and Apple Core AI.