johnmarktaylor91/torchlens

TorchLens – PyTorch 모델의 심층 검사

기능 – TorchLens는 PyTorch 모델의 순전파(및 선택적으로 역전파) 중에 흐르는 모든 텐서를 기록하는 Python 라이브러리입니다. 완전한 계산 그래프 DAG를 구축하고, 방대한 연산자별 메타데이터(형상, 유형, FLOPs, 타이밍, RNG 상태 등)를 첨부하며, 그래프를 PDF/HTML로 시각화할 수 있습니다. 이 캡처는 "증명 가능할 정도로 충실"하며, 라이브러리는 캡처된 DAG를 재생하여 출력이 원래 실행과 일치하는지 확인하고 불일치가 있으면 플래그를 지정합니다.

주요 기능

• 한 줄 추적：log = tl.trace(model, input)는 활성화 값, 그래디언트 및 메타데이터가 포함된 TraceLog 객체를 반환합니다.
• 풍부한 선택기 – tl.func('relu'), tl.in_module('classifier') 또는 사용자 정의 논리 조합과 같은 술어를 사용하여 저장, 중지 또는 개입할 연산자를 선택합니다.
• 순전파 + 역전파 캡처 – 연산자별 그래디언트를 사용할 수 있으며, 그래디언트 저장을 특정 레이어로 제한할 수 있습니다.
• 수용 영역 / 투영 영역 분석 – 경험적 그래디언트 마스크와 교차 검증되는 기하학적 쿼리.
• 개입 – 캡처 중에 제로 어블레이션(zero-ablate), 스케일링, 노이즈 추가, 텐서 교체 등을 즉석에서 수행합니다.
• 시각화 – 자동 생성된 PDF, 접기/펼치기 뷰, 대화형 HTML, Netron 내보내기.
• 데이터셋 전체 추출 – 데이터셋을 일괄 처리하고 활성화 값을 디스크에 기록합니다.
• 성능 고려 – 선택적 캡처는 원시 순전파보다 빠를 수 있으며, 전체 캡처는 약 14배 느리지만 벤치마크가 완료되었습니다.
• 모든 PyTorch 모델에서 작동 – 정적, 동적 제어 흐름, 순환 루프, transformers, GNNs, 확산 모델, MoE 등.

일반적인 워크플로우

import torch, torchvision.models as models, torchlens as tl
model = models.resnet18(weights='IMAGENET1K_V1').eval()
x = torch.rand(1,3,224,224)
log = tl.trace(model, x)          # capture everything
print(log.summary())             # table of layers, FLOPs, params
print(log['conv2d_1_1'].out.shape)  # access a specific activation
log.draw()                       # PDF of the graph

이후 개입, 실행 비교 또는 추적 내보내기가 가능합니다.

설치 – 그래프 그리기를 위해 Graphviz가 필요합니다. pip install torchlens(CPU 전용 코어)로 설치하며, HuggingFace 모델을 지원하기 위해 torchlens[hf]와 같은 옵션을 추가할 수 있습니다.

대상 사용자 – 신경망 내부의 세밀한 검사, 동적 아키텍처 디버깅, 또는 분석, 프로빙, 모델 기반 개입을 위한 체계적인 활성화 값/그래디언트 추출이 필요한 연구자 및 엔지니어.

제한 사항 – 역전파 캡처는 PyTorch에서만 작동합니다. 다른 백엔드는 두 번째 AD 패스를 통해 리프 수준의 그래디언트를 얻습니다. 플랫폼 지원은 주로 Linux/CPU이며, macOS/Windows의 CI 검사는 제한적입니다.

---

링크

• 논문：https://www.nature.com/articles/s41598-023-40807-0
• Model Menagerie (8.5k개의 캡처된 모델 카탈로그)：https://modelmenagerie.ai
• 튜토리얼, 성능 가이드 및 API 문서는 저장소의 docs/ 및 notebooks/에 있습니다.