딥러닝 속도·메모리 최적화의 핵심! Automatic Mixed Precision(AMP) 정리

딥러닝 모델을 학습하다 보면 연산 속도가 너무 느리거나, GPU 메모리가 부족해 학습이 중단되는 경우가 많습니다. 이를 해결할 수 있는 강력한 도구가 바로 Automatic Mixed Precision(AMP)입니다.

이번 포스트에서는 AMP가 무엇인지, 왜 필요한지, 어떻게 사용하는지, 그리고 학습 이후 추론 시 주의할 점까지 한 번에 정리해드릴게요.

---

1. AMP란 무엇인가요?

AMP는 32비트(float32)와 16비트(float16) 연산을 자동으로 섞어서 수행하는 기술입니다. 핵심 목적은 다음 두 가지예요:

• 연산 속도 향상: 16비트 연산은 32비트보다 빠릅니다. 특히 Volta 이후 NVIDIA GPU에서는 Tensor Core 덕분에 학습 속도가 2~4배 빨라질 수 있어요.
• 메모리 절약: float16은 메모리 사용량이 절반이기 때문에, 더 큰 모델 또는 배치를 사용할 수 있습니다.

---

2. 왜 "Automatic"인가요?

수동으로 float16과 float32를 나눠서 코딩하면 너무 복잡하죠. AMP는 프레임워크가 알아서 판단해서 다음과 같이 자동으로 처리합니다:

• 안전한 연산 (예: 행렬곱) → float16으로 처리
• 정밀도가 중요한 연산 (예: Softmax, 손실 계산) → float32로 유지

즉, 사용자는 복잡한 타입 전환을 신경 쓸 필요 없이 성능을 최적화할 수 있습니다.

---

3. AMP는 어떻게 동작하나요?

AMP의 작동 방식은 다음과 같습니다:

1. 모델 파라미터는 float32로 저장
2. 연산 중에는 float16을 사용할 수 있는 부분만 자동 변환
3. Gradient는 float16으로 계산, 최종적으로 float32로 변환하여 파라미터를 업데이트
4. Loss Scaling 기법으로 underflow 문제 방지

---

4. Loss Scaling이란?

float16은 표현 가능한 숫자의 범위가 작기 때문에, 작은 gradient가 0으로 사라지는(underflow) 문제가 발생할 수 있습니다. 이를 방지하기 위해 "Loss Scaling"이라는 기법이 사용됩니다.

• Loss 값에 큰 수를 곱해서 gradient 값을 키운 후 역전파 수행
• Optimizer로 업데이트하기 전에 다시 원래 크기로 나누는 방식

PyTorch의 GradScaler는 이 과정을 자동으로 관리해줍니다.

비유하자면...
작은 목소리(gradient)가 너무 작아 안 들려서 마이크 볼륨(scaling factor)을 키워서 전달하고, 마지막에 다시 줄여서 저장하는 것과 비슷해요.

---

5. PyTorch에서 AMP 쓰는 방법

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for data, target in train_loader:
    optimizer.zero_grad()
    with autocast():  # AMP 적용 구간
        output = model(data)
        loss = loss_fn(output, target)
    scaler.scale(loss).backward()  # 손실 스케일 후 역전파
    scaler.step(optimizer)         # 스케일 해제 후 옵티마이저 스텝
    scaler.update()                # 스케일 팩터 업데이트

• autocast()는 자동으로 연산을 섞어 수행
• GradScaler는 loss scaling 자동 관리

---

6. AMP 학습 이후, 추론에서는 어떻게?

AMP로 학습한 모델은 보통 파라미터는 float32로 저장되어 있습니다. 추론 시엔 다음 두 가지 방법이 있어요:

1. 그대로 float32로 추론 → 가장 안전
2. 모델을 float16으로 변환 후 추론 → 속도와 메모리 절약 가능

model = model.half()        # 파라미터 포함 float16으로 변환
model.eval()
with torch.no_grad():
    output = model(input.half())  # 입력도 float16

---

7. 주의할 점

• 정밀도가 중요한 분야(금융, 과학 계산 등)에서는 float16 사용에 주의
• 일부 연산은 float16 미지원 → 오류 가능
• CPU 추론 시 float16은 오히려 느릴 수 있음 (GPU에서만 이점 큼)

---

8. 핵심 요약

항목	설명
파라미터	항상 float32로 유지
연산	가능하면 float16으로 수행
gradient	float16으로 계산, float32로 변환하여 업데이트
loss scaling	underflow 방지용
추론	float32 또는 float16 변환 가능 (float16은 GPU 권장)

마무리:
Automatic Mixed Precision은 단순한 성능 최적화 기법이 아니라, 현대 딥러닝 모델 학습의 필수 도구입니다. PyTorch, TensorFlow 등 대부분의 프레임워크에서 지원하며, 어렵지 않게 도입할 수 있어요.
AMP를 잘 활용하면 더 빠르고 더 효율적인 학습이 가능해집니다!

---