Stonesjtu/pytorch_memlab

pytorch_memlab – PyTorch용 가벼운 CUDA 메모리 디버깅 도구상자

무엇인가요 – Python만으로 작동하는 패키지로, PyTorch 코드의 각 줄이 얼마나 많은 GPU 메모리를 사용하는지 시각화하고, 장치에 존재하는 텐서를 목록으로 보여주며, 일시적으로 모든 것을 CPU로 이동할 수 있습니다. 명령줄에서도, Jupyter/IPython의 마법 명령어를 통해도 작동합니다.

주요 기능

행 단위 메모리 프로파일러 (LineProfiler / @profile 데코레이터) – line_profiler와 유사하지만, 각 소스 줄당 CUDA의 active_bytes와 reserved_bytes를 보고합니다.
메모리 리포터 (MemReporter) – 모든 활성 torch.Tensor 객체를 나열하고, 기반의 UntypedStorage 크기를 표시할 수 있으며, 공유 가중치, 기울기, 상세한 스토리지 재사용 다이어그램도 처리합니다.
예의 바람직한 모드 – 필요에 따라 모든 GPU 텐서를 CPU로 되돌릴 수 있는 헬퍼 클래스로, 체크포인트를 저장하지 않고 GPU를 해제하는 데 유용합니다.
IPython 통합 – %mlrun과 %%mlrun 마법 명령어를 사용해 한 번의 명령어로 함수나 전체 노트북 셀을 프로파일링할 수 있습니다.
GPU 선택 – set_target_gpu를 통해 실행 중에 프로파일러가 감시할 장치를 전환할 수 있습니다.

설치 방법

# PyPI에서 안정 버전 설치
pip install pytorch_memlab

# GitHub에서 최신 코드 직접 설치
pip install git+https://github.com/stonesjtu/pytorch_memlab

(IPython 마법 명령어를 사용하려면 선택적 추가 패키지를 설치하세요: pip install pytorch_memlab[ipython].)

일반적인 워크플로우

함수 프로파일링
```
from pytorch_memlab import profile

@profile
def train_step(x):
    net = torch.nn.Linear(1024, 1024).cuda()
    return net(x).mean()
```
스크립트가 종료된 후, 각 줄이 얼마나 많은 GPU 메모리를 할당했는지 보여주는 테이블이 출력됩니다.

활성 텐서 확인

from pytorch_memlab import MemReporter
reporter = MemReporter(model)   # model은 선택 사항
reporter.report()               # 간결한 테이블 출력
reporter.report(verbose=True)  # 스토리지 공유 화살표 표시

노트북에서 사용하기

%load_ext pytorch_memlab
%%mlrun -f train_step
train_step(torch.randn(512, 1024).cuda())

GPU 일시 해제

from pytorch_memlab import Courtesy
c = Courtesy()
c.yield_memory()   # 텐서를 CPU로 이동
# …신호를 기다림…
c.restore()        # 다시 이동

왜 중요한가요 – 딥러닝 모델 개발 시 메모리 부족(OOM) 충돌은 흔한 골칫거리입니다. pytorch_memlab는 숨겨진 CUDA 할당 동작을 시각화하여, 비효율적인 텐서 수명 주기, 의도하지 않은 스토리지 공유, 또는 잊혀진 중간 버퍼를 발견하는 데 도움을 줍니다.

제한 사항 / 알려진 문제

프로파일러는 Python이 인식하는 텐서만 볼 수 있습니다. autograd에 의해 생성된 저수준 C 측 버퍼는 "보이지 않는" 것으로 보고되며, 여전히 메모리를 사용할 수 있습니다.
모든 텐서가 CPU에 있더라도, PyTorch CUDA 컨텍스트 자체가 약 1GB의 GPU 메모리를 확보하고 있으며, 현재 이 메모리는 라이브러리로 해제할 수 없습니다.
CUDA 지원 GPU에서만 작동합니다. AMD 또는 CPU 전용 실행에는 지원하지 않습니다.

프로젝트 상태 – 활발히 유지 관리 중 (최신 릴리스 0.3.2, 2026년 6월). CI에서 테스트를 실행하고 PyPI에 웨일을 게시합니다. README에는 사용 예제, 데모 노트북, 변경 기록이 포함되어 있습니다.

위의 모든 정보는 리포지토리의 README에서 직접 가져온 것이며, 외부 주장은 추가되지 않았습니다.

Stonesjtu/pytorch_memlab

pytorch_memlab – PyTorch용 가벼운 CUDA 메모리 디버깅 도구상자

관련