Train ClothingCNN on Fashion-MNIST

curl -LsSf https://astral.sh/uv/install.sh | sh
mkdir -p clothing-android-litert && cd clothing-android-litert
uv init --python 3.11
uv venv
uv add torch==2.13.0 torchvision==0.28.0 --index pytorch-cpu=https://download.pytorch.org/whl/cpu
uv run python -c "import torch, torchvision; print(torch.__version__, torchvision.__version__); print('MPS:', torch.backends.mps.is_available())"
uv run python train.py

Fashion-MNIST is a set of 28×28 grayscale images and ten clothing classes. Keep this exact label order: T-shirt/top, Trouser, Pullover, Dress, Coat, Sandal, Shirt, Sneaker, Bag, Ankle boot. The Android app must use the same order.

Complete training script

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

DEVICE = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
BATCH_SIZE = 128
NUM_CLASSES = 10
transform = transforms.ToTensor()  # float32 grayscale, [0, 1]
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, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(), nn.Linear(64 * 7 * 7, 128),
            nn.ReLU(), nn.Linear(128, NUM_CLASSES),
        )
    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)
print("training on", DEVICE)
for epoch in range(3):
    model.train()
    for images, labels in train_loader:
        images, labels = images.to(DEVICE), labels.to(DEVICE)
        optimizer.zero_grad()
        loss_fn(model(images), labels).backward()
        optimizer.step()
    model.eval()
    correct = total = 0
    with torch.no_grad():
        for images, labels in test_loader:
            predictions = model(images.to(DEVICE)).argmax(1).cpu()
            correct += (predictions == 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")
with torch.no_grad():
    image, label = test_data[0]
    logits = model(image.unsqueeze(0))
    assert image.shape == (1, 28, 28)
    assert logits.shape == (1, NUM_CLASSES)
    print("sample:", label, "->", logits.argmax(1).item())

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 example.dtype == torch.float32
assert logits.shape == (1, 10)
assert logits.dtype == torch.float32
print("contract OK")

clothing.pt stores the trained weights, so keep the ClothingCNN class for conversion. Move the model to CPU before you export. Then copy this file to the Linux machine you use for conversion.