[Pytorch] checkpoint vs torchscript vs onnx 모델 속도 비교

파이토치 모델 변환의 대표적인 3가지 포맷인

체크포인트, 토치스크립트 그리고 onnx 방식에 대하여 inference 과정의 속도를 비교해보는

간단한 실험을 진행해 보겠습니다.(동일한 GPU를 사용하는 환경을 가정해 보겠습니다.)

 

모델 정의 및 변환 코드 예시

우선, 아래와 같이 파라미터 수가 아주 많은 신경망 모델이 하나 있다고 가정해 보겠습니다.

import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()

        self.fc1 = nn.Linear(50000, 20000)
        self.fc2 = nn.Linear(20000, 2000)
        self.fc3 = nn.Linear(2000, 10)

    def forward(self, x):
        return self.fc3(self.fc2(self.fc1(x)))

model = Net()

 

checkpoint, torchscript 및 onnx 파일로 해당 모델을 변환하는 코드 예시는 다음과 같습니다.

# 파이토치 체크포인트로 저장
torch.save(model, 'model.pt')

# torchscript로 저장
model_script = torch.jit.script(model)
model_script.save('model.ts')

# onnx로 저장
dummy_input = torch.randn(1, 50000)
torch.onnx.export(model, dummy_input, 'model.onnx', export_params=True)

 

 

1. checkpoint 포맷 추론 속도

위에서 저장한 체크포인트 파일을 기준으로 10번의 inference 속도를 재보겠습니다.

import time

# 파이토치 체크포인트 모델
ckpt_model = torch.load('model.pt').cuda()
ckpt_model.eval()

for i in range(10):
    dummy_input = torch.randn(1, 50000).cuda()

    start = time.time()
    pred = ckpt_model(dummy_input)
    print(time.time() - start)
    
# 결과
2.761475086212158
0.00020194053649902344
0.0001513957977294922
0.00014710426330566406
0.00014400482177734375
0.00015211105346679688
0.00014925003051757812
0.00014901161193847656
0.00014972686767578125
0.00016355514526367188

첫 추론 시에는 모델을 load하는 과정이 포함되어 시간이 많이 소요되었습니다.

다만 그 이후에는 0.00015초 전후의 매우 빠른 추론 속도가 측정되었습니다.

 

반응형

 

2. torchscript 포맷 추론 속도

이번에는 같은 방식으로 torchscript 파일로 저장된 모델을 불러와 시간을 측정하겠습니다.

# torchscript 모델
ts_model = torch.jit.load('model.ts').cuda()
ts_model.eval()

for i in range(10):
    dummy_input = torch.randn(1, 50000).cuda()

    start = time.time()
    pred = ts_model(dummy_input)
    print(time.time() - start)
    
# 결과
0.02107548713684082
0.05068159103393555
0.00877070426940918
0.0001983642578125
0.0001876354217529297
0.0001628398895263672
0.00017118453979492188
0.00016951560974121094
0.00015497207641601562
0.00021982192993164062

처음 3번의 과정에서는 checkpoint 만큼은 아니지만 시간이 많이 소요되었습니다.

다만, 이후의 과정에서는 0.00018초 전후의 비교적 빠른 속도가 측정되었습니다.

 

 

3. onnx 포맷 추론 속도

마지막으로 같은 방식으로 onnx 파일로 저장된 모델의 추론 속도도 살펴보겠습니다.

import onnxruntime

def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

# onnx 모델
onnx_session = onnxruntime.InferenceSession("model.onnx", providers=['CUDAExecutionProvider'])

for i in range(10):
    dummy_input = torch.randn(1, 50000).cuda()
    ort_inputs = {onnx_session.get_inputs()[0].name: to_numpy(dummy_input)}

    start = time.time()
    pred = onnx_session.run(None, ort_inputs)
    print(time.time() - start)
    
# 결과
0.0302274227142334
0.016111135482788086
0.01609039306640625
0.016085147857666016
0.016072511672973633
0.016064882278442383
0.01606273651123047
0.016077280044555664
0.01609945297241211
0.016098737716674805

처음 추론 과정에서의 모델을 load하는 시간이 상대적으로 적었으나,

이후 과정에서도 0.016초 전후의 비교적 긴 추론 시간이 소요된 것을 볼 수 있었습니다.

 

 

결론

초기의 inference 과정에서는 모델을 불러오는 시간이 포함되어 다소 긴 시간이 소요됩니다.

 

다만, 이후의 과정들을 기준으로 측정된 해당 실험에서의 속도 차이는 다음과 같았습니다.

checkpoint >= torchscript >> onnx

 

다만, 모델의 구조(여기서는 nn.Linear layer만 존재했습니다.),

GPU/CPU 종류 및 운영체제 환경 등에 따라서
실험 결과는 천차만별로 달라질 수 있다는 점을 유의해주시고

참고용으로만 봐주시면 감사하겠습니다.