feat: implement script to export mnist model from pytorch to onnx

Author: Gumichocopengin8

.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+  "python.analysis.extraPaths": ["./pytorch/.venv/lib/python3.12/site-packages/"]
+}

pytorch/.gitignore

@@ -1,3 +1,5 @@
-venv/
+.venv/
 data/
+__pycache__/
 mnist_cnn.pt
+mnist_cnn.onnx

pytorch/README.md

@@ -2,9 +2,13 @@
 
 ## How to Run
 
-- `python3 -m venv venv`
-- `source venv/bin/activate`
+- `python3 -m venv .venv`
+- `source .venv/bin/activate`
 - `pip install -r requirements.txt`
 - `python mnist.py`
 
 The code is from https://github.com/pytorch/examples/tree/main/mnist.
+
+### Export Model from Pytorch to ONNX
+
+- `python export_to_onnx.py`

pytorch/export_to_onnx.py

@@ -0,0 +1,17 @@
+# Convert pt file to onnx file
+import torch
+from mnist import Net
+from pathlib import Path
+
+
+def main():
+    MODEL_PATH = Path("./mnist_cnn.pt")
+    mnist_model = Net()
+    mnist_model.load_state_dict(torch.load(MODEL_PATH))
+    mnist_model.eval()
+    dymmy_input = torch.zeros(1, 1, 28, 28)
+    torch.onnx.export(mnist_model, dymmy_input, "mnist_cnn.onnx", verbose=True)
+
+
+if __name__ == "__main__":
+    main()

pytorch/requirements.txt

@@ -1,2 +1,4 @@
 torch==2.4.0
 torchvision==0.19.0
+onnx==1.16.2
+onnxruntime==1.19.0