torch.cuda.amp GradScaler 공부해보기

이번에 대회를 나가면서 baseline 코드를 수행했는데 코드 일부에 torch.cuda.amp.GradScaler 및 autocast 가 있는 것을 확인하였다. 이전에도 본적이 있어서 알아봤던 거 같은데 기억이 안나 다시 한번 공부하고 기록해놓으려고 한다.

 

우선 이 코드를 왜 쓰는지 알아보기 위해서는 Auto Mixed Precision 이 필요하다. 그리고 이걸 알아보기 전에 부동소수점 즉 Floating Point에 대한 정보를 필요로 한다. 나도 잘 몰라서 기본적으로만 적으려고 한다. 부동소수점은 한가지 방식이 아니라 여러가지 방식으로 나타낼 수 있는데 FP16, FP32, FP64, FP 128이 있다. 값을 세밀하게 표현하기 위해서는 더 높은 숫자인 FP128, FP 64를 사용해야하지만 그만큼 메모리를 더 많이 먹는 단점이 있다. 

 

우리는 주로 FP 32를 사용한다. 그러나 딥러닝 모델이 크면 클수록 그 안에있는 숫자들이 많기 때문에 이 부동소수점으로 인한 메모리가 많이 할당되어 배치크기를 크게 하지 못하고, 결과적으로 학습속도가 저하되는 결과를 낳는다. 그렇다면 가장 기본적으로 생각할 수 있는 방법이 FP32 로 표현한 숫자들을 FP16으로 표현하면 안될까? 그러나 이렇게 다 바꿀 경우

 결과는 위 이미지와 같다. 학습이 잘 수렴하지 않는 것이다. 회색선이 FP16이다. 

 

그렇다면 우리가 쓰는 딥러닝 모델들은 어떤 숫자들로 이루어져있을까? 바로 Activation, Activation Gradient, Weights, Weights Gradients 들이 있다. 그리고 실험을 통해 Weights, Weights Gradient의 숫자들은  FP16으로도 충분히 표현할 수 있었다고 한다. 그렇다면 직관적으로 FP16, FP32를 혼합해서 사용한다면 성능이 변화가 없는 선에서 학습속도를 늘릴 수 있는 것이고 우리가 궁극적으로 궁금했던 GradScaler 및 Autocast가 이러한 역할을 하고 있다는 것이다.

""" define loss scaler for automatic mixed precision """
# Creates a GradScaler once at the beginning of training.
scaler = torch.cuda.amp.GradScaler()

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

  with torch.cuda.amp.autocast():
    # Casts operations to mixed precision 
    outputs = model(inputs)
    loss = criterion(outputs, labels)

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

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

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

실제 코드를 살펴보면 GradScaler를 호출한 다음에 여기서 backward, update해주는 것을 확인할 수 있다. 안에서 어떻게 구동되는지 궁금한 사람은 추가로 알아보면 될 거 같다.