[NVIDIA APEX] Amp에 대해 알아보자 (Automatic Mixed Precision)

version update

• 20-07-25 : amp 모듈이 pytorch 1.5.0 버전부터 기본 라이브러리에 추가되고 있음!

 

pytorch 를 이용해 모델을 학습하다 보면 더 많은 batch size를 학습시키고 싶고 더 빠르게 학습시키고 싶은 생각이 굴뚝같아진다... 하지만 우리가 가지고 있는 데스크탑이나 서버 환경을 물리적으로 확장시키는 방법은 돈이 많이 든다. 돈을 들이지 말고 코드 몇 줄 만으로 모델을 최적화 시키고 batch size를 늘릴 수 없을까?

 

https://github.com/NVIDIA/apex

NVIDIA/apex

NVIDIA에서 'A Pytorch EXtension'(APEX)이라는 패키지를 만들었는데 이 패키지가 어느정도 위 요구를 들어주는 것 같다. (나는 개인적으로 대만족) APEX 패키지에는 mixed precision training과 distributed training 기능이 들어있고 이 글에서는 mixed precision training(amp) 방법에 대해 소개한다.

 

자료 출처 - NVIDIA GTC 2019 talks

Amp - automatic mixed precision

---

apex 깃허브에서 amp를 다음과 같이 소개하고 있다.

 

apex.amp is a tool to enable mixed precision training by changing only 3 lines of your script. 
Users can easily experiment with different pure and mixed precision training modes by supplying different flags to amp.initialize.

(코드상에 3줄 만으로 mixed precision training(혼합 정밀도 학습?)을 할 수 있게 만드는 라이브러리)

 

mixed precision training이란

• 처리 속도를 높이기 위한 FP16(16bit floating point)연산과 정확도 유지를 위한 FP32 연산을 섞어 학습하는 방법
• Tensor Core를 활용한 FP16연산을 이용하면 FP32연산 대비 절반의 메모리 사용량과 8배의 연산 처리량 & 2배의 메모리 처리량 효과가 있다

 

(최적의)사용 조건

• Tensor Core를 이용하여 FP16 연산이 가능한 Volta 이상의 NVIDIA 그래픽 카드 (V100, RTX2080ti, ...)

 

Volta 이전의 그래픽카드더라도 FP16 연산이 가능하다면 amp 기능을 이용할 수 있다. 하지만 성능 효과는 크지 않을 수 있음(관련 블로그 글)

 

얼마나 좋아지는데?

---

속도

- 눈에 띄게 빨라짐

최적화가 되어 배치를 늘릴 수 있기 때문에 학습 속도가 빨라지지만 배치 뿐만 아니라 모델 최적화도 이루어지기 때문에 속도가 증가함

 

배치 크기가 2배로 증가하지만 속도는 2.5배, 3배씩 증가하는 것으로 보아 최적화를 위한 속도 향상도 존재함

성능

- 비슷하거나 조금 향상

 

성능이 향상되는 이유는 배치사이즈에 증가에 따른 학습 효과로 판단됨

 

[사용 방법] #핵간단, #Only 3 lines

---

# Declare model and optimizer as usual, with default (FP32) precision
model = torch.nn.Linear(D_in, D_out).cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

# Allow Amp to perform casts as required by the opt_level
model, optimizer = amp.initialize(model, optimizer, opt_level="O1")
...
# loss.backward() becomes:
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()
...

1. 로드한 model과 optimizer를 amp.initialize로 감싼다
2. 학습 중에 loss와 optimizer를 amp.scale_loss로 감싼다 (as scaled_loss)
3. 감싼 scaled_loss로 back propagation을 진행한다

실제 적용 방법

[실제 적용 방법]

• loss 계산해주고, optimizer에 zero_grad 걸어준 다음
• amp.scale_loss로 감싸서 backward 진행
• optimizer step 진행

 

[핵심 키워드] OPT_LEVEL : 최적화 기법

---

O0

• Only FP32 training
• AMP 적용 안 시킨 것 - baseline accuracy가 된다고 생각하면 됨

 

O1

• 가장 기본적으로 쓰는 OPT_LEVEL
• Tensor Core에 맞는 연산들(ops)은 FP16으로 계산하고 정확한 계산이 필요한 부분은 FP32로 계산하는 Mixed Precision Training
• 별도로 지정하지 않는 한 dynamic loss scaling을 적용하여 학습함
  • Model weights는 FP32로 유지시키기 때문에 loss를 이에 맞게 scaling 해야함

 

O2

• O1 보다 더 FP16을 사용하는 Mixed Precision Training
• batchnorm weights를 제외한 model weights를 FP16으로 casting
• batchnorm, master weights는 FP32로 유지시키기 때문에 역시 dynamic loss scaling이 필요함

 

O3

• Only FP16 training
• 정말 속도만을 위한 모델 최적화
• 만일 모델이 batch normalization을 사용하고 있다면 keep_batchnorm_fp32 = True 로 인자를 전달해줘야 함

 

nvidia apex documents에서는 다음과 같이 설명하고 있음

  O0 and O3는 not true mixed precision. 하지만, 각각 accuracy / speed baseline을 설정할 때 유용하다
  O1 and O2는 mixed precision의 다른 버전. 둘 다 사용해서 최선의 결과를 주는 최적화 방법을 선택하라

 

[Checkpointing] amp로 학습한 모델을 다시 재학습 시킬 때 방법

---

To properly save and load your amp training, we introduce the amp.state_dict(), which contains all loss_scalers and their corresponding unskipped steps, as well as amp.load_state_dict() to restore these attributes.

 

저장할 때 amp.state_dict()를 이용해서 저장하면 amp로 학습했던 때의 파라미터 값들을 불러올 수 있다

# Initialization
opt_level = 'O1'
model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)

# Train your model
...

# Save checkpoint
checkpoint = {
    'model': model.state_dict(),
    'optimizer': optimizer.state_dict(),
    'amp': amp.state_dict()
}
torch.save(checkpoint, 'amp_checkpoint.pt')
...

# Restore
model = ...
optimizer = ...
checkpoint = torch.load('amp_checkpoint.pt')

model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
amp.load_state_dict(checkpoint['amp'])

# Continue training
...

• 가급적이면 다시 실행시킬 때 opt_level은 이전과 같은 것으로 선택하고
• load_state_dict는 amp.initialize 후에 실행해야 한다

 

오류 발견 그리고 torch.cuda.amp 

---

NVIDIA apex amp 모듈을 이용해 모델을 학습시키다가 다음과 같은 상황에서 오류를 발견했다.

• Pytorch의 torch.nn.DataParallel 모듈을 이용하여 multi-gpu 학습을 수행
• apex amp 최적화 기법 O1가 아닌 다른 기법을 이용하여 학습을 수행

 

즉, DataParallel과 amp를 같이 쓰면 오류가 발생할 수 있다... 조심

https://github.com/NVIDIA/apex/issues/227

Data parallel error with O2 and not O1 · Issue #227 · NVIDIA/apex

그렇게 글들을 읽다보니 pytorch 에서 자체적으로 apex 패키지를 반영시켰다는 것을 확인!

 

파이토치 1.5.0 버전부터 amp 모듈이 내부 라이브러리에 들어오기 시작했다. 하지만 아직 1.5.1버전(7월 25일 현재)까지는 dynamic loss scaling에 해당하는 torch.cuda.amp.GradScaler만 사용이 가능하다. 따라서 amp 기능을 온전히 쓰고 싶으면 preview (Nightly)버전을 이용하면 되고 pytorch의 amp 사용 방법은 파이토치 공식 문서에 친절하게 설명되어 있다.

 

https://pytorch.org/docs/master/notes/amp_examples.html

Automatic Mixed Precision examples — PyTorch master documentation

pytorch preview Build 설치 방법 (Linux with Anaconda)

conda install pytorch torchvision cudatoolkit=10.2 -c pytorch-nightly