PyTorch AMP 레시피 — 자동 혼합 정밀도 적용법
PyTorch AMP 레시피 — 자동 혼합 정밀도 적용법
torch.cuda.amp(요즘은 torch.amp)는 일부 연산은 float32, 일부는 float16으로 실행해 성능을 높이는 자동 혼합 정밀도(Automatic Mixed Precision) 를 제공해요.
어느 연산이 빠른가
- 선형·합성곱 같은 연산은
float16(또는bfloat16)에서 Tensor Core 덕에 훨씬 빠르고 메모리도 절약돼요. - 리덕션(reduction) 같은 연산은
float32의 동적 범위가 필요해요. - 서로를 고르는 작업을
torch.autocast컨텍스트가 자동으로 처리해요.
GradScaler의 역할
작은 그래디언트가 float16에서 언더플로(0으로 소실) 하는 걸 막기 위해 gradient scaling 을 해요.
scaler = torch.amp.GradScaler("cuda")
with torch.autocast(device_type="cuda", dtype=torch.float16):
loss = loss_fn(model(x), y)
scaler.scale(loss).backward()
scaler.step(opt)
scaler.update()
Tensor Core 지원 GPU(Volta·Turing·Ampere)에서 2~3배 속도 향상을 볼 수 있고, 학습/추론 재개 시 scaler.state_dict()를 저장해 bitwise 복원도 가능해요.