torch.cuda.amp.GradScaler

2025. 2. 21. 00:02카테고리 없음

torch.cuda.amp.GradScaler는 PyTorch에서 Mixed Precision Training (혼합 정밀도 학습, FP16) 시 그레디언트 스케일링을 담당하는 클래스
이 기능을 사용하면 FP16 연산의 안정성을 높이고 학습 속도를 가속화할 수 있음

 

scaler = torch.cuda.amp.GradScaler(init_scale=2048, growth_interval=1000)

예시 분석

 

📌 1. init_scale=2048

  • 초기 스케일 팩터 (Gradient Scale Factor)를 2048로 설정
  • 기본값은 2.0 ** 16 = 65536이지만, 여기서는 2048로 낮게 설정됨
  • FP16 연산 시 작은 그레디언트 값이 언더플로우(Underflow)되는 문제를 방지
  • 스케일이 너무 작으면 학습이 느려지고, 너무 크면 오버플로우(Overflow) 발생 가능

📌 2. growth_interval=1000

  • 1000번의 업데이트마다 scaler의 크기를 증가시킴
  • 만약 오버플로우가 발생하지 않으면, 1000번 학습 반복 후 scaler가 점진적으로 증가
  • 그레디언트가 너무 작아지는 걸 방지하면서 안정적인 FP16 학습을 보장

왜 GradScaler가 필요한가?

FP16(half precision) 학습을 할 때, **그레디언트 값이 매우 작아 언더플로우(Underflow)**가 발생할 수 있음
이를 방지하기 위해 GradScaler는 그레디언트를 일정 값(스케일)만큼 확대하여 손실값(Gradient)을 더 안정적으로 유지

 

 

언제 GradScaler를 사용해야 할까?

  • NVIDIA GPU에서 torch.cuda.amp.autocast()를 사용할 경우
  • Mixed Precision Training (FP16) 학습을 안정적으로 수행하고 싶을 때
  • 그레디언트 언더플로우(Underflow) 문제를 방지하고 싶을 때
  • 메모리 사용량을 줄이면서 학습 속도를 가속화하고 싶을 때

 

🚀 결론

초기 그레디언트 스케일 값을 2048로 설정하여 안정적인 학습 가능
1000번 업데이트 후 스케일 값을 증가하여 FP16 학습 성능 최적화
FP16 연산 시 발생하는 언더플로우(Underflow) 문제 방지
GPU 메모리를 절약하면서도 빠른 학습 가능