eoinedge commited on
Commit
ab66444
·
verified ·
1 Parent(s): 9c8a39a

TinyDigitNet: files and card

Browse files
README.md ADDED
@@ -0,0 +1,86 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: apache-2.0
3
+ library_name: pytorch
4
+ pipeline_tag: image-classification
5
+ tags:
6
+ - executorch
7
+ - zephyr
8
+ - tinyml
9
+ - mnist
10
+ - cortex-m
11
+ - esp32
12
+ - int8
13
+ datasets:
14
+ - ylecun/mnist
15
+ metrics:
16
+ - accuracy
17
+ model-index:
18
+ - name: TinyDigitNet
19
+ results:
20
+ - task:
21
+ type: image-classification
22
+ name: Handwritten digit recognition
23
+ dataset:
24
+ type: ylecun/mnist
25
+ name: MNIST test set
26
+ metrics:
27
+ - type: accuracy
28
+ value: 0.9746
29
+ name: Top-1 accuracy (fp32, 10,000 test images)
30
+ ---
31
+
32
+ # TinyDigitNet
33
+
34
+ A 14,410-parameter CNN for 28x28 grayscale digits, sized to run under
35
+ [ExecuTorch](https://github.com/pytorch/executorch) on Zephyr RTOS boards with no
36
+ delegate: it has been run on an RP2040 (Cortex-M0+, Arducam Pico4ML-BLE, reading digits
37
+ from the on-board camera), an ESP32-S3 (Seeed XIAO ESP32S3 Sense), an ESP32 (M5Stack ATOM)
38
+ and Cortex-M33 parts (Arduino UNO Q, nRF5340, MCX N947).
39
+
40
+ ```
41
+ conv1 1->8 3x3, ReLU, maxpool 2 28x28 -> 13x13
42
+ conv2 8->16 3x3, ReLU, maxpool 2 13x13 -> 5x5
43
+ fc1 400->32, ReLU
44
+ fc2 32->10
45
+ ```
46
+
47
+ Input: `(1, 1, 28, 28)` float, pixel/255 normalised with mean 0.1307 and std 0.3081
48
+ (the usual MNIST constants; the firmware applies the same transform to camera crops).
49
+ Output: 10 logits.
50
+
51
+ ## Files
52
+
53
+ | File | What |
54
+ |---|---|
55
+ | `tiny_digit_net.py` | model definition (PyTorch) |
56
+ | `tiny_digit_net.pt` | trained state dict, fp32 |
57
+ | `executorch/fp32.pte` | ExecuTorch program, portable fp32 kernels (62.8 KB) |
58
+ | `executorch/int8-cmsis-m0plus.pte`, `-m4.pte`, `-m33.pte`, `-m7.pte` | int8 via the ExecuTorch Cortex-M quantizer and `CortexMPassManager`, one per core (25.4 KB) |
59
+ | `tflite/fp32.tflite`, `tflite/int8.tflite` | the same weights through a Keras twin, for TensorFlow Lite Micro comparisons (post-training int8 quantization) |
60
+ | `test_vectors.h` | one test image plus reference logits for an on-device self-test |
61
+
62
+ ## Results
63
+
64
+ | Metric | Value |
65
+ |---|---|
66
+ | MNIST test accuracy (fp32) | 97.46 % |
67
+ | Parameters | 14,410 |
68
+ | Inference, RP2040 @ 133 MHz, fp32 portable kernels | 766 ms |
69
+
70
+ ## Training
71
+
72
+ Standard MNIST, Adam, a few epochs on CPU; the script is `train/train_mnist.py` in the
73
+ [Pico4ML-BLE repo](https://github.com/eoinjordan/executorch-zephyr-arducam-pico4ML-BLE).
74
+ Export with `torch.export` then `to_edge_transform_and_lower(...).to_executorch()`; the
75
+ int8 variants use `CortexMQuantizer` + `prepare_pt2e`/`convert_pt2e` and the Cortex-M pass
76
+ manager from ExecuTorch 1.5.
77
+
78
+ ## Where it runs
79
+
80
+ - [executorch-zephyr-arducam-pico4ML-BLE](https://github.com/eoinjordan/executorch-zephyr-arducam-pico4ML-BLE): camera to screen digit reader
81
+ - [executorch-zephyr-xiao-esp32s3](https://github.com/eoinjordan/executorch-zephyr-xiao-esp32s3): same on the XIAO ESP32-S3 Sense
82
+ - [executorch-vs-tflm-benchmarks](https://github.com/eoinjordan/executorch-vs-tflm-benchmarks): ExecuTorch vs TFLM, fp32 and int8, across boards
83
+
84
+ ## License
85
+
86
+ Apache-2.0. MNIST is the public dataset by LeCun, Cortes and Burges.
executorch/fp32.pte ADDED
Binary file (62.8 kB). View file
 
executorch/int8-cmsis-m0plus.pte ADDED
Binary file (25.4 kB). View file
 
executorch/int8-cmsis-m33.pte ADDED
Binary file (25.4 kB). View file
 
executorch/int8-cmsis-m4.pte ADDED
Binary file (25.4 kB). View file
 
executorch/int8-cmsis-m7.pte ADDED
Binary file (25.4 kB). View file
 
test_vectors.h ADDED
@@ -0,0 +1,106 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ /* Generated by models/export.py; do not edit. */
2
+ #pragma once
3
+ #define TEST_INPUT_NUMEL 784
4
+ #define TEST_OUTPUT_NUMEL 10
5
+ #define TEST_OUTPUT_ARGMAX 7
6
+ static const float kTestInput[TEST_INPUT_NUMEL] = {
7
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
8
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
9
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
10
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
11
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
12
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
13
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
14
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
15
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
16
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
17
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
18
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
19
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
20
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
21
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
22
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
23
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
24
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
25
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
26
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
27
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
28
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
29
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
30
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
31
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
32
+ -0.4242130f, -0.4242130f, 0.6449587f, 1.9305104f, 1.5995764f, 1.4977505f, 0.3394811f, 0.0340035f,
33
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
34
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
35
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, 2.4014549f, 2.8087585f,
36
+ 2.8087585f, 2.8087585f, 2.8087585f, 2.6432915f, 2.0959773f, 2.0959773f, 2.0959773f, 2.0959773f,
37
+ 2.0959773f, 2.0959773f, 2.0959773f, 2.0959773f, 1.7395868f, 0.2376552f, -0.4242130f, -0.4242130f,
38
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
39
+ -0.4242130f, -0.4242130f, 0.4285787f, 1.0268056f, 0.4922199f, 1.0268056f, 1.6504892f, 2.4650962f,
40
+ 2.8087585f, 2.4396398f, 2.8087585f, 2.8087585f, 2.8087585f, 2.7578456f, 2.4905527f, 2.8087585f,
41
+ 2.8087585f, 1.3577399f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
42
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
43
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.2078330f, 0.4158505f, -0.2460177f, 0.4285787f, 0.4285787f,
44
+ 0.4285787f, 0.3267528f, -0.1569200f, 2.5796502f, 2.8087585f, 0.9249798f, -0.4242130f, -0.4242130f,
45
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
46
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
47
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, 0.6322305f, 2.7960303f,
48
+ 2.2359879f, -0.1951047f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
49
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
50
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
51
+ -0.4242130f, -0.1441918f, 2.5414655f, 2.8214867f, 0.6322305f, -0.4242130f, -0.4242130f, -0.4242130f,
52
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
53
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
54
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, 1.2177293f, 2.8087585f, 2.6051068f,
55
+ 0.1358293f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
56
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
57
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
58
+ 0.3267528f, 2.7451172f, 2.8087585f, 0.3649375f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
59
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
60
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
61
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, 1.2686423f, 2.8087585f, 1.9559668f, -0.3605718f,
62
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
63
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
64
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.3096589f,
65
+ 2.1850750f, 2.7323890f, 0.3140246f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
66
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
67
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
68
+ -0.4242130f, -0.4242130f, -0.4242130f, 1.1795444f, 2.8087585f, 1.8923256f, -0.4242130f, -0.4242130f,
69
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
70
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
71
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, 0.5304046f, 2.7705739f,
72
+ 2.6305633f, 0.3012964f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
73
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
74
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
75
+ -0.4242130f, -0.1823765f, 2.3887267f, 2.8087585f, 1.6886740f, -0.4242130f, -0.4242130f, -0.4242130f,
76
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
77
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
78
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.3860283f, 2.1596186f, 2.8087585f, 2.3632703f,
79
+ 0.0212752f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
80
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
81
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
82
+ 0.0594599f, 2.8087585f, 2.8087585f, 0.5558611f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
83
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
84
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
85
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.0296377f, 2.4269114f, 2.8087585f, 1.0395339f, -0.4114847f,
86
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
87
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
88
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, 1.2686423f,
89
+ 2.8087585f, 2.8087585f, 0.2376552f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
90
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
91
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
92
+ -0.4242130f, -0.4242130f, 0.3522093f, 2.6560197f, 2.8087585f, 2.8087585f, 0.2376552f, -0.4242130f,
93
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
94
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
95
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, 1.1159033f, 2.8087585f,
96
+ 2.8087585f, 2.3632703f, 0.0849164f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
97
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
98
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
99
+ -0.4242130f, -0.4242130f, 1.1159033f, 2.8087585f, 2.2105315f, -0.1951047f, -0.4242130f, -0.4242130f,
100
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
101
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
102
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
103
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
104
+ -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f, -0.4242130f,
105
+ };
106
+ static const float kTestOutput[TEST_OUTPUT_NUMEL] = {-1.2525061f, -2.3244209f, 4.0803161f, 2.5801816f, -5.6759157f, -5.8199277f, -7.9415760f, 10.5989647f, -2.2308075f, -0.1515399f};
tflite/fp32.tflite ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3113668584b82da7759699c13fbf9071b6477fe83fc8d3b8cab7f841fcaac1dd
3
+ size 60988
tflite/int8.tflite ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:f1fbf6622d8e85a9cc7086acfcede4a466a6ea1901996fc9822e20b2c9ba596a
3
+ size 19992
tiny_digit_net.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:06b7e67c3239fbe4a60fe241f59fd43eb6a4ace5131188558794851623dc98e0
3
+ size 61111
tiny_digit_net.py ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) 2026 Eoin Jordan
2
+ # SPDX-License-Identifier: Apache-2.0
3
+ """TinyDigitNet: a ~14k parameter CNN for 28x28 grayscale digits.
4
+
5
+ Kept deliberately small so the fp32 portable ExecuTorch kernels run it in a
6
+ few milliseconds on the ESP32-S3 (Xtensa LX7, 240 MHz) with no delegate.
7
+ """
8
+
9
+ import torch
10
+ import torch.nn as nn
11
+ import torch.nn.functional as F
12
+
13
+ # MNIST normalisation constants. The firmware applies the same transform to
14
+ # camera frames, so keep these in sync with firmware/src/main.cpp.
15
+ MEAN = 0.1307
16
+ STD = 0.3081
17
+
18
+ INPUT_SHAPE = (1, 1, 28, 28)
19
+ NUM_CLASSES = 10
20
+
21
+
22
+ class TinyDigitNet(nn.Module):
23
+ def __init__(self) -> None:
24
+ super().__init__()
25
+ self.conv1 = nn.Conv2d(1, 8, kernel_size=3) # 28 -> 26
26
+ self.conv2 = nn.Conv2d(8, 16, kernel_size=3) # 13 -> 11
27
+ self.fc1 = nn.Linear(16 * 5 * 5, 32)
28
+ self.fc2 = nn.Linear(32, NUM_CLASSES)
29
+
30
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
31
+ x = F.max_pool2d(F.relu(self.conv1(x)), 2) # 26 -> 13
32
+ x = F.max_pool2d(F.relu(self.conv2(x)), 2) # 11 -> 5
33
+ x = x.flatten(1)
34
+ x = F.relu(self.fc1(x))
35
+ return self.fc2(x)