https://github.com/feifeibear/pytorchmemtracer

Depict GPU memory footprint during DNN training of PyTorch
https://github.com/feifeibear/pytorchmemtracer

dnn memory oom pytorch

Last synced: 3 months ago
JSON representation

Depict GPU memory footprint during DNN training of PyTorch

• Host: GitHub
• URL: https://github.com/feifeibear/pytorchmemtracer
• Owner: feifeibear
• Created: 2022-01-24T06:03:53.000Z (almost 4 years ago)
• Default Branch: master
• Last Pushed: 2022-11-17T07:17:55.000Z (about 3 years ago)
• Last Synced: 2025-06-11T20:52:44.259Z (7 months ago)
• Topics: dnn, memory, oom, pytorch
• Language: Python
• Homepage:
• Size: 120 KB
• Stars: 11
• Watchers: 1
• Forks: 2
• Open Issues: 1
• Metadata Files:
  • Readme: README.md

Awesome Lists containing this project

README

          
# A Memory Tracer For PyTorch

OOM is a nightmare for PyTorch users.

However, most of them do not know the exact memory footprint during training.

This project helps you depict the GPU memory usage changing curve during training.

We record the peak GPU memory of an operator at the moment of the operator finished.

The operator can be computing of a FWD submodule or a BWD submodule.

## Usage

```

python train.py 

python visualize.py memstats.pkl

```

![alt perf](./memstats.png "an example result")