autocast: Mixed Precision

Mixed Precision

Pytorch는 기본적으로 모델 가중치를 32-bit float를 사용한다. 모델에 입력하기 전 .float()를 사용하는 것도 이 때문이다. 그런데 Mixed Precision은 필요에 따라 16-bit float를 섞어 사용한다. 이러한 방식은 여러 장점이 있다.

• 32-bit가 아닌 16-bit를 사용해 GPU 메모리 사용량을 줄인다.
• 학습 속도가 빨라진다.
• 학습 성능이 비슷하거나 약간 향상된다. (NVIDIA)

물론 16-bit는 32-bit보다 표현할 수 있는 범위가 좁아 정밀도가 떨어진다. 따라서 값이 underflow/overflow 되지 않도록 scaling 해야 한다. 

pytorch는 위 과정을 아주 쉽게 적용할 수 있다.

---

autocast

from torch import autocast
from torch.amp import GradScaler

scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()

        # autocast를 활용한 Forward pass.
        with autocast(device_type="cuda", dtype=torch.float16):
            output = model(input)
            loss = loss_fn(output, target)
        
        # loss 값 조정 & backward pass
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        # scaler 조정
        scaler.update()

• autocast: 자동으로 일부 값을 16-bit로 조정
• autocast 안에서 forward pass
• GradScaler를 이용해 loss 범위 조정
• autocast 밖에서 backward pass
• scaler를 이용해 step optimizer
• 다음 epoch을 위한 scaler update

scaler는 scaler factor를 이용해 loss 값이 너무 작아지거나 너무 커지는 것을 방지한다. 이 scaler factor는 loss에 따라 자동으로 조정된다.