[Model Inference] Pytorch 2.0 Compile 사용 후기 및 장단점 | pytorch compile 모델 추론 속도 개선 테스트

Pytorch Compile 튜토리얼

Pytorch 2.0 Overview

 

Pytorch 2.0

Speedups for torch.compile against eager mode on an NVIDIA A100 GPU

 

compiled_model = torch.compile(model)

긴 설명할 것 없이 Pytorch 2.0 이후 compile 이라는 것이 추가되었는데, 위 예시처럼 torch.comile(model)이라는 짧은 코드 한 줄만 추가하면 모델 인퍼런스 속도를 향상시킬 수 있다고 한다.

 

공식 문서에 따르면 A100 GPU에서 모델 학습 속도는 43% 향상되고, 모델 인퍼런스 속도는 Float32 precision에서 21%, AMP precision에서 51% 정도 향상된다고 한다.

 

 

Pytorch 2.0 Compile 모델 인퍼런스 테스트

torchvision에서 제공하는 기본적인 모델인 resnet50으로 pytorch compile의 속도 개선 테스트를 진행해봤다.

 

# 테스트 코드

import torch
import torchvision.models
import numpy as np

model = torchvision.models.resnet50()
model.eval().cuda()

starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
repetitions = 10
timings=np.zeros((repetitions,1))

B = 1
H = 1000
W = 1000
input_data = torch.randn((B, 3, H, W)).float().cuda()

for _ in range(5):
    _ = model(input_data)
torch.cuda.synchronize()

with torch.no_grad():
    for rep in range(repetitions):
        starter.record()
        torch_out = model(input_data)
        ender.record()
        torch.cuda.synchronize()
        curr_time = starter.elapsed_time(ender)
        timings[rep] = curr_time
        # new_H = torch.randint(H-100,H+100,(1,))
        # new_W = torch.randint(W-100,W+100,(1,))
        # input_data = torch.randn((1, 3, new_H, new_W)).float().cuda()

print('torch 모델 평균 소요 시간 : ', np.mean(np.array(timings)))

# model_compiled = torch.compile(model, dynamic=True)
model_compiled = torch.compile(model)
del model

for _ in range(5):
    _ = model_compiled(input_data)
torch.cuda.synchronize()


timings=np.zeros((repetitions,1))
with torch.no_grad():
    for rep in range(repetitions):
        starter.record()
        complied_out = model_compiled(input_data)
        ender.record()
        torch.cuda.synchronize()
        curr_time = starter.elapsed_time(ender)
        timings[rep] = curr_time
        # new_H = torch.randint(H-100,H+100,(1,))
        # new_W = torch.randint(W-100,W+100,(1,))
        # input_data = torch.randn((1, 3, new_H, new_W)).float().cuda()

print('compiled 모델 평균 소요 시간 : ', np.mean(np.array(timings)))

# 결과 비교
error = torch.abs(torch_out - complied_out).mean()
print(f"Mean Absolute Error: {error.item()}")

• 테스트 모델 : resnet50
• 테스트 방법
  • pytorch 모델과 compile한 모델을 생성한 후 몇 가지 입력 크기에 대해 속도 테스트
  • 입력 텐서를 모델에 10회 반복해서 입력한 후 평균 소요 시간을 측정
• 입력 크기
  • 정적인 입력 크기로 테스트
  • torch.compile()의 파라미터에서 dynamic=True로 설정하면 dynamic input shape에 대응이 가능하다 하여 테스트

 

 

# 테스트 결과

Input shape	인퍼런스 평균 소요 시간 (ms)		속도 향상 (%)
	Pytorch model	Compiled model
[1,3,500,500]	9.78	14.56	-32.83
[10,3,500,500]	44.78	38.15	17.38
[1,3,1000,1000]	20.46	17.88	14.43
[8,3,1000,1000]	140.84	113.8	23.76
[1,3,900~1100,900~1100] (Dynamic shape)	162.26	8513	-98.09

• 왜인지 [1,3,500,500] 크기의 작은 입력 크기에서는 compile 모델이 오히려 속도 감소
• 입력 크기는 작더라도 배치 크기를 키우면 속도 향상 
• 입력 크기를 키우면 배치=1 에서도 속도 향상되고, 배치가 커질수록 속도 향상이 두드러짐
  • 전반적으로 15~20%의 속도 향상 (공식 문서와 비슷한 수치)
• 동적인 입력에서의 테스트는... 속도가 오히려 매우 느려짐
  • 공식 문서에 따르면 현재는 Danamic shape 입력에 대한 지원이 제한되어 있다고 함

→ Pytorch Compile 사용 시 사용하는 모델과 입력 크기에 따라 속도 개선 정도가 다를 수 있으니 테스트를 한 후에 사용하는 것이 좋을 것 같다.

→ 학습 시에는 배치가 입력 크기가 고정된 경우가 많으니 상관없지만, 인퍼런스 시 동적인 입력을 사용하는 경우에는 사용이 제한될 수 있다.