PyTorch amp 문서 — autocast와 GradScaler
PyTorch amp 문서 — autocast와 GradScaler
torch.amp 는 PyTorch의 혼합 정밀도 학습·추론을 위한 네임스페이스예요. 두 핵심 구성 요소가 있어요.
torch.autocast
autocast 컨텍스트 내부의 CUDA 연산을 텐서 코어에 유리한 정밀도(기본은 float16, torch.bfloat16도 가능)로 자동 캐스팅해요. 연산 종류에 따라 더 나은 정밀도를 고르는 오토캐스트 연산 참조를 제공해요. 추론·평가 전방 패스에도 단독으로 쓰기에 좋아요.
GradScaler
계산 그래프를 건드리지 않고 그래디언트를 스케일(곱)·언스케일하는 도구예요.
scaler.scale(loss).backward()로 스케일된 손실로 역전파해 언더플로를 막아요.scaler.step(opt)가 그래디언트에 inf/NaN이 없으면 옵티마이저를 갱신하고, 있으면 건너뛰어요.scaler.update()로 다음 반복의 스케일을 조정해요.
조언
손실이 inf/NaN이면 캐스팅 범위의 연산 영역을 float32로 강제하거나, 자세한 예제 문서(amp_examples)를 참고해요. 커스텀 C++ 연산은 디스패처 autocast와 함께 다뤄요.