PyTorch 자동 혼합 정밀도(Automatic Mixed Precision, AMP) 가이드

주제: PyTorch의 torch.cuda.amp 모듈을 활용한 혼합 정밀도 학습(Automatic Mixed Precision)으로, 메모리 절약 및 연산 속도 향상을 실현하는 방법을 정리합니다.

---

1. AMP란 무엇인가?

• 딥러닝 학습 시, 기본적으로 32비트 부동소수점(FP32) 연산을 사용합니다.
• 반정밀도(FP16) 연산은 메모리 사용량을 절반으로 줄이고, 특정 GPU에서 더 빠른 처리 속도를 제공합니다.
• Automatic Mixed Precision(AMP)은 모델 파라미터는 FP32로 유지하면서 연산(곱셈·덧셈 등)만 FP16으로 자동 전환하여, 정밀도 손실 없이 메모리와 연산 효율을 극대화하는 기술입니다.

---

2. 왜 사용해야 할까?

이슈	FP32만 사용 시	AMP 적용 시
메모리 사용량	높음 (모든 텐서가 FP32)	절반 수준 (텐서·활성화가 FP16)
연산 속도	표준 GPU 성능 활용	Tensor Core 사용 시 2×~3×↑
코드 복잡도	없음	간단한 컨텍스트 관리자 추가
수치 안정성	최고	스케일러 사용으로 안정성 보장

---

3. PyTorch AMP 사용 예시

import torch
from torch.cuda.amp import autocast, GradScaler

# 1) 모델, 옵티마이저, 스케일러 초기화
model = MyModel().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scaler = GradScaler()  # loss 스케일러

# 2) 학습 루프
for inputs, targets in dataloader:
    inputs, targets = inputs.cuda(), targets.cuda()

    optimizer.zero_grad()

    # 3) autocast 컨텍스트: FP16/FP32 자동 전환
    with autocast():
        outputs = model(inputs)
        loss = criterion(outputs, targets)

    # 4) 스케일링된 역전파
    scaler.scale(loss).backward()
    # 5) 옵티마이저 스텝
    scaler.step(optimizer)
    # 6) 스케일러 업데이트
    scaler.update()

1. autocast() 블록 안의 연산은 자동으로 FP16 또는 FP32로 전환됩니다.
2. GradScaler는 작은 loss 값에서 언더플로우를 방지하기 위해 loss를 스케일 업 후 역전파하고, 스케일을 동적으로 조절합니다.

---

4. 장단점 비교

구분	장점	단점
메모리	활성화·파라미터 메모리 최대 50% 절감	스케일러 오버헤드 소폭 증가
속도	Tensor Core 활용 시 2×~3× 학습 속도 향상	FP16 지원 하드웨어 필요
구현 난이도	autocast, GradScaler 추가만으로 적용 가능	일부 연산(FP16 비지원 연산) 주의
수치 안정성	스케일러로 언더/오버플로우 최소화	매우 작은 loss 함수에서 불안정 가능

---

5. 최적화 팁

• FP16 지원 여부 확인: Ampere 이상 아키텍처 GPU에서 Tensor Core 활용 가능
• 연산 예외 처리: autocast() 내에서 FP16 지원이 미흡한 연산(Softmax 등)을 FP32로 강제 처리 가능
• 스케일 초기값 조정: GradScaler(init_scale=2**16) 등 하이퍼파라미터 튜닝
• 로그 모니터링: torch.cuda.amp.GradScaler.get_scale()으로 스케일 변화 추적

---

🙌 이제 PyTorch AMP로 메모리와 성능을 동시에 잡아보세요!