mixed precision

출처:Gpt

원래 float32를 쓰는데, 계산 정확도가 엄밀하지 않아도 되는 부분에는 float 16을 사용하여 연산 효율화 하는 방법. (메모리 절약, 속도 향상)

 

모델 훈련하는 부분에 적용해주면 되는데, 

아래 autocast부분에 모델 예측 + loss계산은 float16로,

optimizing하는 부분은 float32로

scaler = torch.cuda.amp.GradScaler()

for inputs, labels in dataloader:
    optimizer.zero_grad()
    with torch.cuda.amp.autocast():
        outputs = model(inputs)
        loss = criterion(outputs, labels)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

- gradscaler 

loss: 0.5 → backward 계산 → gradient: 1e-8 (float16에서는 표현 불가) → gradient = 0!

scaler.scale(loss).backward()   # ⬅️ 이때는 scaled된 loss 사용
scaler.step(optimizer)          # ⬅️ 여기서 optimizer에게 float32 복구된 gradient 전달
scaler.update()                 # ⬅️ 다음 step을 위한 scale 값 자동 조정