[DL] Mixed Precision Training: 효율적인 딥러닝을 위한 접근 방법

Mixed Precision Training 소개

Mixed Precision Training은 다양한 부동소수점 데이터 유형을 혼합하여 사용하는 기법입니다. 주로 float32 대신 float16을 사용하여 메모리 사용량을 줄이고 학습 속도를 향상시킵니다. 이 방법은 최신 GPU에서 특히 효과적입니다.

Mixed Precision Training의 장점

• 메모리 절약: Float16은 float32에 비해 절반의 메모리만을 사용합니다. 이를 통해 더 큰 모델이나 미니 배치로 학습이 가능해집니다.
• 속도 향상: 최신 GPU는 float16 연산에 최적화되어 있어, 학습 속도가 상당히 향상될 수 있습니다.
• 정규화 효과: 수치적 불안정성이 약간의 정확도 향상으로 이어질 수 있습니다.

주의 사항

• 수치적 불안정성: Float16은 제한된 표현 범위로 인해 underflow나 overflow 문제가 발생할 수 있습니다.
• 불완전한 하드웨어 지원: 일부 하드웨어는 mixed precision을 완전히 지원하지 않을 수 있습니다.

AMP (Automatic Mixed Precision)

PyTorch에서는 torch.cuda.amp 패키지를 통해 AMP를 제공합니다. 이에는 두 가지 주요 구성 요소가 포함됩니다: Autocast와 GradScaler.

Autocast

• 기능: 연산의 정밀도를 동적으로 조절합니다. autocast 컨텍스트 관리자 내에서 실행되는 연산은 자동으로 float16으로 실행됩니다.
• 예시:

with torch.cuda.amp.autocast():
    output = model(input)
    loss = criterion(output, target)

GradScaler

• 기능: float16 연산 시 발생할 수 있는 gradient underflow나 overflow 문제를 완화합니다. Loss의 스케일을 조절하여 gradient 값을 관리합니다.
• 예시:

scaler = torch.cuda.amp.GradScaler()

with torch.cuda.amp.autocast():
    output = model(input)
    loss = criterion(output, target)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

결론

Mixed Precision Training은 학습 속도를 향상시키고 메모리 사용량을 줄이는 데 유용한 기법입니다. autocast와 GradScaler는 AMP를 구현하는 데 중요한 도구로, 연산의 정밀도를 동적으로 관리하고 gradient의 안정성을 유지합니다. 이 기법을 사용할 때는 수치적 불안정성과 하드웨어 호환성에 주의해야 합니다.