Mixed Precision - fp16으로 학습하기

딥러닝 학습을 할 때 실수를 표현하는 방법으로 FP32(single precision)를 사용한다.
FP 뒤 32는 32bit를 사용하여 실수를 저장한다는 의미이다.
따라서 FP16(Half Precision)을 사용한다면 GPU를 더 효율적으로 활용하고 연산 속도도 빨라진다.
하지만 당연히 정밀도가 떨어져 오차들이 누적되면 제대로 정답에 수렴하지 못할 수도 있다.
이 한계를 극복하기 위해 나온 Mixed Precision Training을 활용하여 더 빠르고 효율적인 학습을 해보았다.

사용방법

https://github.com/hoya012/automatic-mixed-precision-tutorials-pytorch
위 깃허브 주소에 가면 잘 정리되어 있다.

Before (기존)

for batch_idx, (inputs, labels) in enumerate(data_loader):
  self.optimizer.zero_grad()

  outputs = self.model(inputs)
  loss = self.criterion(outputs, labels)

  loss.backward()
  self.optimizer.step()

After (fp16)

""" define loss scaler for automatic mixed precision """
scaler = torch.cuda.amp.GradScaler()

for batch_idx, (inputs, labels) in enumerate(data_loader):
  self.optimizer.zero_grad()

  with torch.cuda.amp.autocast():
    outputs = self.model(inputs)
    loss = self.criterion(outputs, labels)

  # Scales the loss, and calls backward() 
  # to create scaled gradients 
  self.scaler.scale(loss).backward()

  # Unscales gradients and calls 
  # or skips optimizer.step() 
  self.scaler.step(self.optimizer)

  # Updates the scale for next iteration 
  self.scaler.update()