PyTorch 2.0에서 달라지는 점 - torch.compile

PyTorch 2.0

 

출처: 파이토치 공식 블로그

 

PyTorch 2.0은 22년 12월 PyTorch Conference에서 발표되었고, 23년 3월 정식 릴리즈 되었다. 이전의 PyTorch 1.x 버전들보다 빠르고, Pythonic하고 Dynamic하다고 한다. 어떤 점들이 달라졌을지 한번 알아봅시다.

 

 

torch.compile

torch.compile은 PyTorch 2.0의 메인 API이다. 모델을 미리 컴파일하여 속도를 높이는 기술이다. torch.compile은 TorchDynamo, AOTAutograd, PrimTorch, TorchInductor 네 가지의 새로운 기술을 기반으로 만들어졌다. 각 기술에 대한 자세한 설명은 여기에서 찾아볼 수 있다.

 

torch.compile이 나오게 된 배경

 

사용법

 

torch.compile은 기존의 모델에 한 줄만 추가하면 사용할 수 있다.

 

compiled_model = torch.compile(model)

 

이렇게 컴파일된 모델은 기존과 동일하게 사용할 수 있다. 다음은 ResNet18에 torch.compile을 적용하는 예시 코드이다. torch.compile()을 하는 부분을 제외하면 기존의 PyTorch 모델과 사용법이 동일한 것을 확인할 수 있다.

 

import torch
import torchvision.models as models

model = models.resnet18().cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
compiled_model = torch.compile(model)

x = torch.randn(16, 3, 224, 224).cuda()
optimizer.zero_grad()
out = compiled_model(x)
out.sum().backward()
optimizer.step()

 

처음 compiled_model을 실행할 때에는 모델 컴파일에 시간이 걸리지만, 이후 실행들은 더 빨라진다.

 

 

한가지 주의해야 할 점은, 컴파일된 모델을 저장할 때 state_dict만 저장할 수 있다. 즉,

 

torch.save(optimized_model.state_dict(), "foo.pt")

 

이건 되고

 

torch.save(optimized_model, "foo.pt")

 

이건 안된다.

 

 

속도 비교

 

torch.compile을 사용했을 때의 속도를 비교하기 위해 PyTorch는 총 163가지의 오픈소스 모델에 대한 실험을 진행했다. 실험에 사용한 모델은 다음과 같다.

• HuggingFace Transformer 모델 46가지
• TIMM 모델 61가지
• TorchBench 모델 56가지

 

모델을 전혀 수정하지 않고 코드에 torch.compile만 추가했을 때, 훈련 속도가 43% 빨라졌다고 한다 (NVIDIA A100 GPU 기준). 한편 NVIDIA 3090과 같은 데스크탑용 GPU를 사용했을 때는 속도 향상이 그보다는 덜하다고 한다.

 

 

Example

 

아래 블로그에서 torch.compile을 이용한 실험을 진행하고, 속도를 비교했다.

A Quick PyTorch 2 Tutorial - Zero to Mastery Learn PyTorch for Deep Learning

 

NVIDIA TITAN RTX GPU를 이용했고, CIFAT10 데이터셋에 ResNet50 모델로 실험했다.

 

 

 

왼쪽은 epoch=5, 오른쪽은 epoch=15에 대한 실험 결과이다. (multiple run은 5epoch 훈련을 여러번 실행시켰다는 뜻)

 

Epoch 수가 늘어날수록 컴파일된 모델의 속도 향상이 더 큰 것을 확인할 수 있다. 이는 첫 컴파일에 시간이 오래 걸리기 때문이다.

속도 차이가 크지 않은 것은 데스크탑용 GPU를 사용했기 때문이고, A100등의 산업용 GPU를 사용하면 속도향상이 더 크다고 한다.