| # Copyright (c) 2026 Eoin Jordan | |
| # SPDX-License-Identifier: Apache-2.0 | |
| """TinyDigitNet: a ~14k parameter CNN for 28x28 grayscale digits. | |
| Kept deliberately small so the fp32 portable ExecuTorch kernels run it in a | |
| few milliseconds on the ESP32-S3 (Xtensa LX7, 240 MHz) with no delegate. | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| # MNIST normalisation constants. The firmware applies the same transform to | |
| # camera frames, so keep these in sync with firmware/src/main.cpp. | |
| MEAN = 0.1307 | |
| STD = 0.3081 | |
| INPUT_SHAPE = (1, 1, 28, 28) | |
| NUM_CLASSES = 10 | |
| class TinyDigitNet(nn.Module): | |
| def __init__(self) -> None: | |
| super().__init__() | |
| self.conv1 = nn.Conv2d(1, 8, kernel_size=3) # 28 -> 26 | |
| self.conv2 = nn.Conv2d(8, 16, kernel_size=3) # 13 -> 11 | |
| self.fc1 = nn.Linear(16 * 5 * 5, 32) | |
| self.fc2 = nn.Linear(32, NUM_CLASSES) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| x = F.max_pool2d(F.relu(self.conv1(x)), 2) # 26 -> 13 | |
| x = F.max_pool2d(F.relu(self.conv2(x)), 2) # 11 -> 5 | |
| x = x.flatten(1) | |
| x = F.relu(self.fc1(x)) | |
| return self.fc2(x) | |