gpu_memory_profiling

Profile the GPU memory usage of every line in a pytorch code

Example Usage

python example_mnist.py

Dependency

This code depends on py3nvml. Pip install is available here:

pip install py3nvml

Tested with pytorch version 0.4.0 and py3nvml verion 0.2.0

Acknowledgement

The gpu_profile.py is a modified version of this script.

The example_mnist.py is modified from this script.