PyTorch 2.0 Tutorial

이번에 PyTorch 2.0이 발표되었다. 아직 한국어로는 많은 리뷰가 없는 것 같아서, 쉽게 따라 할 수 있는 Tutorial을 작성하고자 한다.

 

이번 PyTorch 2.0의 핵심 아이디어는 바로 컴파일러의 지원이다. 많은 사람들이 알고 있겠지만 PyTorch는 이전까지 컴파일러를 사용하지 않았다. 이번에 컴파일러가 추가되면서 엄청난 속도 향상을 이뤄냈다.

 

대부분의 사람들은 컴파일러 엔지니어가 아니기 때문에 간단히 설명하도록 하겠다. 모르겠으면 그런가보다 하고 넘어가도 괜찮다. (사실상 User는 사용 방법만 알면 된다.)

 

전체적인 Compiler의 Overview이다.

 

 

Dynamo

제일 먼저 그림 위에 있는 Dynamo는 Frontend에 해당하는데, 흔히 우리가 사용하는 Python Script를 읽어와서 FX graph를 만들게 된다. (FX graph는 여기서 중요하지 않으니 궁금한 사람은 따로 찾아보기 바란다.)

 

또한, 인공지능 엔지니어면 알 듯이, PyTorch는 Backward 계산을 자동으로 해주는데, AOTAutograd에서 Backward에 대한 FX graph를 얻을 수 있다.

 

여기서 얻을 수 있는 것이 Aten/Prims IR이다. IR(Intermediate Representation)은 그냥 컴파일러에서 Machine Code를 만들기 전에 최적화를 위해 사용되는 중간 표현이라고 생각하면 된다.

 

 

Inductor

Inductor는 Backend에 해당된다. 앞에서 얻은 Aten or Prims IR을 Loop Level IR로 내린다. 그리고 이를 Machine-specific code로 Mapping한다.

 

여기서 Inductor는 C++에 사용되는 OpenMP나, OpenAI에서 출시한 Triton을 사용한다.

 

Triton 또한 하나의 큰 주제기 때문에 더 깊이 알고 싶다면 따로 공부해 보는 것을 추천한다.

 

이렇게 Backend를 마치게 되면 최적화된 code generation이 된다.

 

정리하면 위의 그림과 같다.

 

Tutorial Code

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()

 

사용법은 그냥 우리가 사용하던 Model 혹은 함수에 torch.compile() 만 추가하면 된다.

 

 

참고로 모델을 저장할 때 주의할 점이 있다.

 

torch.save(optimized_model.state_dict(), "foo.pt")
# both these lines of code do the same thing
torch.save(model.state_dict(), "foo.pt")

 

state_dict()를 저장하는 것은 compile 된 모델을 바로 저장해도 문제가 없다.

 

torch.save(optimized_model, "foo.pt") # Error
torch.save(model, "foo.pt")           # Works

 

하지만 모델 object 자체를 바로 저장하는 것은 아직 지원하지 않는다. 컴파일되지 않은 모델은 바로 저장이 가능하다.

 

 

 

아직은 초기 버전이기 때문에 작은 오류들이 많이 발생하는 것 같다. 점차 지나면서 많이 해결될 것으로 보인다. 혹시 관심이 있다면 한번 써보는 것을 추천한다.

 

 

참고 자료: PyTorch Documents