Create and train a PyTorch model for digit classification using the MNIST dataset
Introduction
Prepare a PyTorch Development Environment
Create a PyTorch model for MNIST
About PyTorch Model Training
Perform Training and Save the Model
Deploy the Model for Inference
Learn about Inference on Android
Create an Android Application
Prepare the Test Data
Run the Application
Optimizing Neural Network Models in PyTorch
Create an optimized PyTorch model for MNIST
Run optimization
Update the Android application
Next Steps

  Log an issue

  Fork and edit

  Discuss on Discord

About this Learning Path

Skill level:  Advanced
Reading time:  2 hrs 40 min
Last updated:  03 Aug 2026
Skill level:
  Advanced
Reading time:
  2 hrs 40 min
Last updated:
  03 Aug 2026
Author:Dawid Borycki
Arm IP:Cortex-A Neoverse
Tags: ML Windows Linux macOS Android Studio Visual Studio Code
Author:
  • Dawid Borycki
Arm IP:
Cortex-A Neoverse
Tags:
ML Windows Linux macOS Android Studio Visual Studio Code

Who is this for?

This is an advanced topic for software developers interested in learning how to use PyTorch to create and train a feedforward neural network for digit classification, and also software developers interested in learning how to use and apply optimizations to the trained model in an Android application.

What will you learn?

Upon completion of this Learning Path, you will be able to:

  • Prepare a PyTorch development environment.
  • Download and prepare the MNIST dataset.
  • Create and train a neural network architecture using PyTorch.
  • Create an Android app and load the pre-trained model.
  • Prepare an input dataset.
  • Measure the inference time.
  • Optimize a neural network architecture using quantization and fusing.
  • Deploy an optimized model in an Android application.

Prerequisites

Before starting, you will need the following:

  • A machine that can run Python3, Visual Studio Code, and Android Studio.
  • For the OS, you can use Windows, Linux, or macOS.

Summary

AI-assisted

This summary was drafted with an approved AI-assisted workflow and reviewed by Arm contributors before publication. Human technical review remains part of the process so the final page reflects engineering rigor, accuracy, and Arm editorial standards.

Close
?
You’ll build and train a PyTorch feedforward neural network for MNIST digit classification, then use it for inference and Android deployment. You’ll set up Python, prepare the dataset, and train the model. Then, you’ll reload the model’s saved parameters and apply the same preprocessing to new images. You’ll quantize and fuse an optimized variant, integrate it into an Android application, and compare inference times before deployment.

Frequently asked questions

AI-assisted

These FAQs were drafted with an approved AI-assisted workflow and reviewed by Arm contributors before publication. Human technical review remains part of the process so the final page reflects engineering rigor, accuracy, and Arm editorial standards.

Close
?
How do I know the MNIST dataset is set up correctly before training?
Confirm that the training and test splits download without errors and that the DataLoader returns batches. With batch_size set to 32, the first batch contains 32 images of 28x28 pixels and corresponding labels from 0–9.
What should I look for during training to confirm the model is learning?
Monitor the loss and confirm that it decreases over epochs. Evaluate on the test data periodically to check whether predictions align with the true labels more often.
After training, what do I need to load the model for inference?
Use the model file produced during training and provide its path when loading. Recreate the same model architecture before loading the saved parameters.
How do I validate that inference works on new images?
Apply the training preprocessing, including normalization and tensor conversion, then run a prediction. The output maps to a digit from 0–9. Compare it with a known test-set label.
How do I compare unoptimized and optimized models after quantization and fusing?
Measure inference time with the original model, then repeat the measurement after quantization and fusing with identical inputs and conditions. Use the timings to choose the model for the Android application.
Next