This repo contains the code for trainig a digit classification model in PyTorch based on the MNIST dataset, and converting the trained model to the ONNX format.
MNIST_Pytorch.py
- Create and train the Pytorch model, and save the state dictionary of the trained model.PyTorch_Model.pt
- The saved state dictionary of the PyTorch model.PyTorch_To_ONNX.py
- Contains code to convert the PyTorch model to ONNX format.ONNX_Model.onnx
- The final ONNX format model.