This script uses PyTorch to train a simple neural network on the MNIST dataset. It performs the following steps:
- Downloads the MNIST dataset.
- Loads the entire dataset into GPU memory.
- Trains a model for 10 epochs.
- Collects and saves performance stats (time, accuracy, GPU usage, memory) to the
stats/directory in the project root.
- An NVIDIA GPU with CUDA.
- Python and Pip installed.
- Install dependencies:
pip install torch psutil requests numpy
From the python/ directory, run:
python python.py