AI 모델을 개발하다보면 학습 속도가 너무 느리거나, 컴퓨팅 용량 문제로 CUDA error가 발생하는 것을 경험하게 된다. 이럴 때 모델을 경량화해서 용량 문제와 속도 문제를 모두 해결할 수 있는 방법이 있어 소개하려 한다.
AMP란?
AMP는 Automatic Mixed Precision(자동 혼합 정밀도)의 약자로, 딥러닝 모델 학습 시 FP16(16비트 부동소수점)과 FP32(32비트 부동소수점) 연산을 혼합하여 사용하도록 지원하는 기능입니다. PyTorch와 같은 딥러닝 프레임워크에서 제공하며, 학습 속도를 높이고 메모리 사용량을 줄이면서도 모델의 정확도를 유지할 수 있도록 돕습니다.
쉽게 말하면 pytorch는 모델 학습할 때 FP32(32비트 부동 소수점) 연산으로 파라미터, 가중치 등을 연산한다. 32비트 부동 소수점은 소수점 이하의 값을 32비트까지 포함할 수 있는 자료형이고, torch.float에 해당하는 자료형이다.
반대로 FP16(16비트 부동 소수점)은 소수점 이하의 값을 16비트까지 저장할 수 있는 자료형이기 때문에 연산 시에 FP32보다 정밀도를 떨어지지만, 저장용량을 줄이고 속도를 높일 수 있는 대안이 된다. 그리고 AMP는 딥러닝 모델 학습 과정에서 FP16을 활요할 수 있도록 지원하는 기능이다.
pytorch에서 AMP를 활용하는 방법
AMP는 Pytorch의 torch.cuda.amp모듈이 있어 쉽게 구현이 가능하다. 보통 시간이 가장 오래 소요되는 model의 연산이 이뤄질 때 FP16을 활용한다.
-
GradScaler 초기화
torch.cuda.amp.GradScaler는 손실 값을 스케일링하여 FP16 연산에서 발생할 수 있는 언더플로우 문제를 방지합니다. -
autocast 컨텍스트 매니저 사용
torch.cuda.amp.autocast를 사용하면, 적합한 연산은 FP16으로 처리하고 나머지는 FP32로 자동 처리됩니다.
예시
import torch
from torch.cuda.amp import GradScaler, autocast
# 모델, 옵티마이저, 손실 함수 정의
model = YourModel().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = torch.nn.CrossEntropyLoss()
# AMP 초기화
scaler = GradScaler(enable = True) #True일 때 FP16을 활용함
for epoch in range(num_epochs):
model.train()
for inputs, targets in dataloader:
inputs, targets = inputs.cuda(), targets.cuda()
optimizer.zero_grad()
# Forward pass 과정에서 AMP를 통해 FP 16 활용
with autocast():
outputs = model(inputs)
loss = criterion(outputs, targets)
# 손실 스케일링 후 backward pass
scaler.scale(loss).backward()
# 그래디언트 언스케일링 후 옵티마이저 스텝
scaler.step(optimizer)
scaler.update()주의사항
AMP는 자료형을 바꿔주기 때문에 Scaler를 잘 활용해야 오차 없이 모델의 학습이 이뤄진다. 그렇기 때문에 혼동하지 않고, 원하는 효과를 얻기 위해서는 주의사항을 주의해야한다.
-
수치 불안정성
FP32보다 소숫점 아래 자리수가 적기 때문에 overflow나 underflow가 발생할 수 있다. -
손실 스케일링
FP16으로 loss를 연산했기 때문에 backward process나 optimizer를 업데이트 할 때는 scale을 잘 맞춰줘야한다. -
Loss가 0이 되어 학습이 되지 않는 문제
현재 AMP를 활용해 모델을 학습하던 중에, Loss가 Nan 값으로 반환되는 경우가 있었다. 이 부분을 해결하기 위해 AMP를 빼는 것이 가장 효과적인 방법이지만, AMP를 빼면 CUDA 에러가 발생해 AMP를 최대한 건드리지 않고 제대로된 학습이 될 수 있는 상황을 만드는 것이 필요하다.
결론
AMP를 활용하면 일부 민감도는 떨어질 수 있으나, 결국 속도와 용량 사이에서 큰 변화를 가져오기 때문에 모델 학습과정에서 활용하면 유용한 방법론이 될 것으로 기대된다.
