It trains the LeNet-5 neural network on the MNIST dataset from scratch, and afterwards you can give it newly hand-written digits 0 to 9 and it will hopefully recognize the digit for you.
Install on Ubuntu 24.10 with:We use our own
sudo apt install protobuf-compiler
git clone https://github.com/activatedgeek/LeNet-5
cd LeNet-5
git checkout 95b55a838f9d90536fd3b303cede12cf8b5da47f
virtualenv -p python3 .venv
. .venv/bin/activate
pip install \
Pillow==6.2.0 \
numpy==1.24.2 \
onnx==1.13.1 \
torch==2.0.0 \
torchvision==0.15.1 \
visdom==0.2.4 \
;
pip install
because their requirements.txt uses >=
instead of ==
making it random if things will work or not.On Ubuntu 22.10 it was instead:
pip install
Pillow==6.2.0 \
numpy==1.26.4 \
onnx==1.17.0 torch==2.6.0 \
torchvision==0.21.0 \
visdom==0.2.4 \
;
Then run with:This script:It throws a billion exceptions because we didn't start the Visdom server, but everything works nevertheless, we just don't get a visualization of the training.
python run.py
- does a fixed 15 epochs on the training data
- it then uses the trained net from memory to check accuracy with the test data
- then it also produces a
lenet.onnx
ONNX file which contains the trained network, nice!
The terminal outputs lines such as:
Train - Epoch 1, Batch: 0, Loss: 2.311587
Train - Epoch 1, Batch: 10, Loss: 2.067062
Train - Epoch 1, Batch: 20, Loss: 0.959845
...
Train - Epoch 1, Batch: 230, Loss: 0.071796
Test Avg. Loss: 0.000112, Accuracy: 0.967500
...
Train - Epoch 15, Batch: 230, Loss: 0.010040
Test Avg. Loss: 0.000038, Accuracy: 0.989300
One of the benefits of the ONNX output is that we can nicely visualize the neural network on Netron:
output = self.c1(img)
x = self.c2_1(output)
output = self.c2_2(output)
output += x
output = self.c3(output)