This project fine-tune torchvision pretrained weights to classify MNIST images and try to run Pytorch model with Rust
- Python 3.8.10
- torch 2.0.1
- torchvision 0.15.2
Install requirements
pip install -U pip
pip install -r requirements.txt
You can download MNIST dataset from here, use torch dataset or simply run this script
bash scripts/download_data.sh
Modify config in ./configs/config.py
Simply run
python tools/pytorch/train.py
Some experiment results
Model | Accuracy | Confusion Matrix | Pretrained | Model size |
---|---|---|---|---|
Resnet18 | 99.64% | ![]() |
Model | 44.71MB |
You can download weight file above
You can infer with
python tools/pytorch/infer.py
Refer here
You can try on your own 😉