Mixed Precision Training 이해 및 설명

1. Mixed Precision Training 이란?

대부분의 deep learning framework(e.g. PyTorch, TensorFlow)들은 모델을 training할 때 float32(FP32) data type을 사용하게 됩니다. 즉, 모델의 weight와 input data가 모두 FP32(32bit)의 data type을 가진다는 뜻입니다. 이와 다르게 Mixed-precision training은 single-precision(FP32)와 half-precision(FP16) format을 결합하여 사용하여 모델을 training하는 방식입니다. 

(FP16 data type은 FP32와 다르게 16bit만을 사용하게 됩니다.)

 

Mixed-precision training방식을 통해 다음과 같은 장점을 가집니다.

• FP32로만 training한 경우와 같은 accuracy 성능을 도출
• Training time이 줄음
• Memory 사용량이 줄음
  • 이로 인해 더 큰 batch size, model, input을 사용 가능하게 함

Mixed-precision training은 NVIDIA에 의해 처음 제안되었고 Automatic Mixed Precision(AMP)라는 feature를 개발하였습니다. AMP feature는 특정 GPU operation을 FP32에서 mixed precision으로 자동으로 바꿔주었으며 이는 performance를 향상시키면서 accuracy는 유지하는 효과를 가져왔습니다. 

 

PyTorch 1.6부터는 PyTorch 안에 AMP package(torch.cuda.amp)를 추가하였습니다. 아래와 같이 torch.cuda.amp는 기존의 NVIDIAd의 AMP의 pain point를 보완하였습니다.

• Window OS  지원
• DataParallel과 intra-process model parallelism 지원 (Multi-GPU 지원)
• Gradient penalty (double backward) 지원
  • e.g) L1 regularization, L2 regularization, ...
• sparse gradient 지원

그리고 PyTorch에서 FP16과 FP32으로 autocast가능한 CUDA operation은 아래와 같습니다. 보시면 모델의 구성 layer(e.g. conv, linear, LSTMCell, ...)들은 FP16으로 auto cast가능하신것을 알 수 있습니다.

https://pytorch.org/docs/stable/amp.html#cuda-ops-that-can-autocast-to-float16

2. Mixed Precision Training 사용 예제

※ 저는 PyTorch framework에서만 사용하는 예제를 설명드립니다.

import torch

scaler = torch.cuda.amp.GradScaler() # Training시에 생성

for data, label in data_iter:
   optimizer.zero_grad()
   with torch.cuda.amp.autocast(): # Mixed precision으로 operation들을 casting 
      outputs = model(data)

   scaler.scale(loss).backward() # Loss를 scaling한 후에 backward진행
   scaler.step(optimizer) # 원래 scale에 맞추어 gradient를 unscale하고 optimizer를 통한 gradient update
   scaler.update() # 다음 iteration을 위해 scale update

 

위에서 mixed precision은 오직 with torch.cuda.amp.autocast(): context안에서만 일어나게 되고 이를 통해 위에서 말씀드린 performance효과를 보게 됩니다. 그럼 scaler라는 class instance는 어떤 역할을 하는 것일까요?

 

2.1 Scaler의 역할

Forward-pass시에 FP16으로 계산된 결과를 통해 backward-pass에서도 FP16으로 계산됩니다. 이 때 gradient의 값이 작아 float16으로 표현할 수 없다면 underflow가 발생하게 됩니다. 이런 문제점을 해결하기 위해 loss를 scale factor(이는 GradScaler()를 초기화할때 설정가능)와 곱하여 loss의 값을 크게 만드는 scaler.scale(loss).backward()를 사용하게 됩니다. loss의 값이 커지게 되면 자연스럽게 underflow문제점이 사라지겠죠. 

 

Loss가 scale factor와 곱해졌었는데 weight update전에 원래 scale로 돌려놓기 위해 위의 scale factor만큼 나누어 주어 unscale하게 되는 과정은 scaler.step(optimizer)에서 진행됩니다. 이때, 이미 scaled gradient가 inf, NaN값을 가지면 optimizer.step() 함수는 skip되고 해당 gradients는 weight update에 사용되지 않고 버리게 됩니다.

 

그리고 scaler.update()를 통해 다음 iteration을 위한 scale factor를 업데이트합니다.

3. Mixed precision training 실험

PyTorch framework안에서 mixed precision training을 사용했을 경우와 아닌 경우의 time cost를 비교하려고 합니다. cifar10 dataset을 사용하여 ResNet18, 50, 101, MobileNetv2의 성능을 비교해보겠습니다. 

실험에 사용한 코드는 해당 url 에서 사용가능합니다. 

 

아래 표는 한 epoch당 소요되는 training/test 시간(단위 seconds)을 나타낸것입니다. 사용된 GPU는 V100 1개입니다.

	ResNet18	ResNet50	ResNet101	MobileNetv2
w/o mixed precision	train: 20.4s test: 1.4s	train: 58.0s test: 3.6s	train: 97.1s test: 5.8s	train: 20.4s test: 1.5s
w mixed precision	train: 10.4s test: 1.4s	train: 25.9s test: 3.6s	train: 43.9s test: 5.8s	train: 22.2s test: 1.5s

 

위의 결과를 통해 mixed precision training을 사용하였을 때 time cost가 적게 드는 것을 확인가능합니다.