Train ClothingCNN on Fashion-MNIST

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 train.py

This course runs on a Mac. Train with Apple Silicon MPS when it is available. Otherwise use the CPU. Do not use CUDA.

Fashion-MNIST contains 28×28 grayscale pictures and ten classes. We keep the image prep simple: ToTensor() turns pixels into float32 values in the 0–1 range, and the app will do the same. The class order is fixed. It must stay the same after conversion.

These are the labels: T-shirt/top, Trouser, Pullover, Dress, Coat, Sandal, Shirt, Sneaker, Bag, and Ankle boot.

Complete training script

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# Mac path: prefer MPS on Apple silicon, else CPU. Never CUDA for this course.
if torch.backends.mps.is_available():
    DEVICE = torch.device("mps")
else:
    DEVICE = torch.device("cpu")
print("training on", DEVICE)
BATCH_SIZE = 128
NUM_CLASSES = 10

# ToTensor gives float32 [0, 1], matching the later Swift preprocessing.
transform = transforms.ToTensor()
train_data = datasets.FashionMNIST("data", train=True, download=True, transform=transform)
test_data = datasets.FashionMNIST("data", train=False, download=True, transform=transform)
train_loader = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True)
test_loader = DataLoader(test_data, batch_size=BATCH_SIZE)

class ClothingCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),              # 28x28 -> 14x14
            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),              # 14x14 -> 7x7
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(64 * 7 * 7, 128),
            nn.ReLU(),
            nn.Linear(128, NUM_CLASSES), # logits, no softmax here
        )

    def forward(self, image):
        return self.classifier(self.features(image))

model = ClothingCNN().to(DEVICE)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# A few epochs is enough for a conversion demo. Train longer for a better app.
for epoch in range(3):
    model.train()
    for images, labels in train_loader:
        images, labels = images.to(DEVICE), labels.to(DEVICE)
        optimizer.zero_grad()
        logits = model(images)
        loss = loss_fn(logits, labels)
        loss.backward()
        optimizer.step()

    model.eval()
    correct = total = 0
    with torch.no_grad():
        for images, labels in test_loader:
            logits = model(images.to(DEVICE))
            correct += (logits.argmax(1).cpu() == labels).sum().item()
            total += labels.numel()
    print(f"epoch {epoch + 1}: accuracy={correct / total:.3f}")

model = model.cpu().eval()
torch.save(model.state_dict(), "clothing.pt")

# Check one real Fashion-MNIST sample and its output shape.
with torch.no_grad():
    sample_image, sample_label = test_data[0]
    sample_logits = model(sample_image.unsqueeze(0))
    sample_prediction = sample_logits.argmax(1).item()
    print("sample:", sample_label, "->", sample_prediction)
    assert sample_logits.shape == (1, NUM_CLASSES)

The saved clothing.pt file holds weights, not the Python class. Keep ClothingCNN available when you load the weights for export. The printed sample uses the same dataset prep as training.

Export and conversion run on CPU. That is why the script moves the trained model back to CPU before saving. Training can use MPS. The .aimodel path does not need a CUDA machine.

Check the tensor contract

model = ClothingCNN()
model.load_state_dict(torch.load("clothing.pt", map_location="cpu", weights_only=True))
model.eval()

example = torch.zeros(1, 1, 28, 28, dtype=torch.float32)
logits = model(example)
assert example.shape == (1, 1, 28, 28)
assert logits.shape == (1, 10)

The model accepts NCHW tensors. That means batch, channel, height, and width, in that order. It returns logits (raw scores). Apply softmax only when you want a confidence number to show in the app.

Next, export this exact ClothingCNN and convert it to clothing.aimodel.