Export ClothingCNN for LiteRT · PyTorch to Android with LiteRT

Environment setup

Convert on Linux. The official LiteRT Torch converter currently says it supports Linux. Python 3.11 is recommended. Copy clothing.pt and this script from the Mac to the Linux host.

mkdir -p clothing-android-convert && cd clothing-android-convert
uv init --python 3.11
uv venv
uv add torch torchvision litert-torch
uv run python -c "import torch, litert_torch; print(torch.__version__, 'litert_torch OK')"
uv run python convert.py

If your package resolver cannot install the current wheel, use the project’s documented Linux container or CI image. Do not switch to a converter that is not verified. Do not add CUDA just for this CPU conversion.

This course uses one path: PyTorch module → litert_torch.convert → clothing.tflite. LiteRT Torch uses PyTorch export inside, and it writes the classic TFLite FlatBuffer that LiteRT reads. You do not need an ONNX-to-TensorFlow step in the middle.

Complete conversion script

import torch
from torch import nn
import litert_torch

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, 10),
        )
    def forward(self, image):
        return self.classifier(self.features(image))

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),)
with torch.no_grad():
    edge_model = litert_torch.convert(model, example)
edge_model.export("clothing.tflite")
print("wrote clothing.tflite")

The input is NCHW [1, 1, 28, 28], not NHWC. The output is [1, 10] float32 logits. LiteRT Android buffers are positional, meaning the first buffer is buffer zero. The app writes the image into input buffer zero and reads logits from output buffer zero. Inspect the FlatBuffer before you ship.

Inspect and copy the artifact

ls -lh clothing.tflite
mkdir -p app/src/main/assets
cp clothing.tflite app/src/main/assets/clothing.tflite

The required check is a non-empty file plus a quick Android runtime smoke test. If conversion says an operator is not supported, keep this model architecture and update the pinned LiteRT Torch package on Linux. Do not change the tensor layout as a guess.

NHWC versus NCHW

Do not transpose unless you mean to. Training and export use NCHW float32. The Kotlin code writes row-major pixels into the single NCHW channel as index = y * 28 + x. If a later model’s inspected input is NHWC, change both the export example and the Android buffer writer at the same time.

Copy clothing.tflite into the Android project. The next lesson creates the Gradle app and packages it into the APK.

Key concepts

Convert on Linux with Python 3.11 and uv add torch torchvision litert-torch.
Call litert_torch.convert on model.eval() with a CPU float32 [1, 1, 28, 28] example.
Export clothing.tflite with positional input buffer 0 and output buffer 0.
Copy the file to app/src/main/assets/clothing.tflite.

Takeaways

Do not transpose to NHWC unless you change both export and Kotlin together.
Required check: a non-empty file plus an Android runtime smoke test.
If an operator fails, keep the architecture and update the pinned litert-torch package.