딥러닝 모델 배포하기 #02 - TorchScript & Pytorch JIT

딥러닝 모델 배포하기 시리즈 2편입니다 :)

1편 (MLOps PipeLine과 연산 최적화 / 모델 경량화)을 먼저 읽으시길 권장합니다 😊

TorchScript & Pytorch JIT

지난 포스팅에서 간단하게 Pytorch와 Tensorflow에 대해 설명하였다. link

Production 분야에서 Pytorch는 Tensorflow에 비해 약세를 띄고 있는데, Facebook(Meta)가 tensorflow를 따라잡고자 내놓은 것이 TorchScript와 Pytorch JIT이다. 이들은 Pytorch model을 최적화하여 model serving을 보다 잘할 수 있도록 도와준다.

보통 모델을 production화 하려면, 두가지가 필요하다.

1. Portability
   • 모델이 다양한 환경에서 export 될 수 있어야 함
   • Python interpreter process 에서뿐만이 아니라 C++ server나 mobile /embedded device 에서도 작동이 가능해야함
2. Performance
   • inference latency와 throughput, 모두의 성능을 유지하면서도 최적화를 해야함

Pytorch 는 Python의 특징을 많이 가지고 있는 프레임워크이다. 때문에 Portability와 Performance, 이 두가지 측면에서 약세를 보였고, 이를 해결하기 위해 Torchscript는 코드를 Eager mode에서 Script mode로 변환한다.

Tools to Transition from Eager to Script

• Eager Mode: normal python runtime mode로 prototyping, training, experimenting을 위해 사용된다
• Script Mode
  • production deployment를 위해 변환한 모드
  • runtime 과정에서 Python Interpreter로 실행되지 않기 때문에 병렬 연산, 최적화 등이 가능해진다.

그렇다면 Eager mode에서 Script mode로 어떻게 변환할까?

Pytorch model을 TorchScript로 변환하는 방법에는 두가지 방법이 있다. (1) Tracing 방식과 (2) Annotation 방식이다.

Tracing

Eager To Script Mode with torch.jit.trace()

Tracing 방법은 어떤 입력값(data instance)을 사용하여 모델의 구조를 파악하고, 이 입력값의 모델 안에서의 흐름을 통해 모델을 기록하는 방식이다. 조건문을 많이 사용하지 않는 모델의 경우 이 방식을 이용하여 변환하는 것이 적합하다.

보통 Pytorch 모델을 Tracing을 통해 Torchscript로 변환하려면, 모델의 instance를 예시 input값과 함께 torch.jit.trace 함수에 넘겨주어야한다.

import torch
import torchvision

# An instance of your model.
model = torchvision.models.resnet18()

# An example input you would normally provide to your model's forward() method.
example = torch.rand(1, 3, 224, 224)

# Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
traced_script_module = torch.jit.trace(model, example)

trace된 ScriptModule은 일반적인 Pytorch module과 같은 방식으로 입력값을 받아 처리할 수 있다.

In[1]: output = traced_script_module(torch.ones(1, 3, 224, 224))
In[2]: output[0, :5]
Out[2]: tensor([-0.2698, -0.0381,  0.4023, -0.3010, -0.0448], grad_fn=<SliceBackward>)

Tracing 방식은 eager model의 코드를 재사용할 수 있는 효과적인 방법이다. 그러나 이 방식을 사용하면 Control-flow나 data structure, python construct가 보존되지 않는다. 따라서 이 방식을 사용하는 경우에는 항상 IR을 검사하여 pytorch model이 올바르게 동작하는지를 확인해줘야한다.

이러한 limitation을 해결하기 위해 Annotation 방식이 고안되었다.

Annotation

Eager To Script Mode with torch.jit.script()

예를 들어 다음과 같은 pytorch model이 있다고 가정해보자.

import torch

class MyModule(torch.nn.Module):
    def __init__(self, N, M):
        super(MyModule, self).__init__()
        self.weight = torch.nn.Parameter(torch.rand(N, M))

    def forward(self, input):
        if input.sum() > 0:
          output = self.weight.mv(input)
        else:
          output = self.weight + input
        return output

MyModule model은 input값에 따라 영향을 받는 Control-flow 를 사용하고 있기 때문에 tracing 기법은 적합하지 않다. 대신 torch.jit.script()함수를 통해 모듈을 compile하여 ScriptModule로 변환한다.

또한, 이 방식은 tracing mode와 다르게 data sample은 전달할 필요가 없다. 오직 model의 instance만 input으로 넣어주면 된다.

my_module = MyModule(10,20)
sm = torch.jit.script(my_module)

---

Mixing Tracing and Scripting

다음과 같이 Tracing과 Scripting을 함께 사용할 수도 있다.

Example (calling a traced function in script)

import torch

def foo(x, y):
    return 2 * x + y

traced_foo = torch.jit.trace(foo, (torch.rand(3), torch.rand(3)))

@torch.jit.script
def bar(x):
    return traced_foo(x, x)

Example (calling a script function in a traced function)

import torch

def foo(x, y):
    return 2 * x + y

traced_foo = torch.jit.trace(foo, (torch.rand(3), torch.rand(3)))

@torch.jit.script
def bar(x):
    return traced_foo(x, x)

참고할 만한 자료 ✍🏻

• PyTorch JIT and TorchScript
• TorchScript and PyTorch JIT / Deep Dive
• LOADING A TORCHSCRIPT MODEL IN C++

---

Pytorch code를 torchscript로 변환하는 건 간단하다.

특히 huggingface에서 제공하는 많은 모델들은 pretrained model을 불러오는 과정에서 script mode를 함께 지원하기 때문에 더욱더 간단하다.

• Exporting transformers models (torchscripts)

Example 1: BERT

Part 1

• Initializes BERT Tokenizer & creates sample data

from transformers import BertTokenizer, BertModel
import numpy as np
import torch
from time import perf_counter

def timer(f,*args):   
    
    start = perf_counter()
    f(*args)
    return (1000 * (perf_counter() - start))
    
script_tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', torchscript=True)

# Tokenizing input text
text = "[CLS] Who was Jim Henson ? [SEP] Jim Henson was a puppeteer [SEP]"
tokenized_text = script_tokenizer.tokenize(text)

# Masking one of the input tokens
masked_index = 8
tokenized_text[masked_index] = '[MASK]'
indexed_tokens = script_tokenizer.convert_tokens_to_ids(tokenized_text)
segments_ids = [0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1]

# Creating a dummy input
tokens_tensor = torch.tensor([indexed_tokens])
segments_tensors = torch.tensor([segments_ids])

Part 2

• 2.1 Normal Pytorch Model

# Example 1.1 BERT on CPU
native_model = BertModel.from_pretrained("bert-base-uncased")
np.mean([timer(native_model,tokens_tensor,segments_tensors) for _ in range(100)])

# Example 1.2 BERT on GPU
# Both sample data model need be on the GPU device for the inference to take place
native_gpu = native_model.cuda()
tokens_tensor_gpu = tokens_tensor.cuda()
segments_tensors_gpu = segments_tensors.cuda()
np.mean([timer(native_gpu,tokens_tensor_gpu,segments_tensors_gpu) for _ in range(100)])

• 2.1 TorchScript Model
  • torchscript=True

script_model = BertModel.from_pretrained("bert-base-uncased", torchscript=True)

# Example 2.1 torch.jit.trace on CPU
traced_model = torch.jit.trace(script_model, [tokens_tensor, segments_tensors])
np.mean([timer(traced_model,tokens_tensor,segments_tensors) for _ in range(100)])

# Example 2.2 torch.jit.trace on GPU
traced_model_gpu = torch.jit.trace(script_model.cuda(), [tokens_tensor.cuda(), segments_tensors.cuda()])
np.mean([timer(traced_model_gpu,tokens_tensor.cuda(),segments_tensors.cuda()) for _ in range(100)])

• Runtime 시간 비교

TorchScript로 Pytorch model을 변환하면, Pytorch JIT에 의해 빠르게 Compile 될 수 있기 때문에 코드를 실행시키는 것이 훨씬 빨라진다.

---

Example 2: ResNet

import torchvision
import torch
from time import perf_counter
import numpy as np

def timer(f,*args):   
    start = perf_counter()
    f(*args)
    return (1000 * (perf_counter() - start))
  
# Example 1.1 Pytorch cpu version

model_ft = torchvision.models.resnet18(pretrained=True)
model_ft.eval()
x_ft = torch.rand(1,3, 224,224)
np.mean([timer(model_ft,x_ft) for _ in range(10)])

# Example 1.2 Pytorch gpu version

model_ft_gpu = torchvision.models.resnet18(pretrained=True).cuda()
x_ft_gpu = x_ft.cuda()
model_ft_gpu.eval()
np.mean([timer(model_ft_gpu,x_ft_gpu) for _ in range(10)])

# Example 2.1 torch.jit.script cpu version

script_cell = torch.jit.script(model_ft, (x_ft))
np.mean([timer(script_cell,x_ft) for _ in range(10)])

# Example 2.2 torch.jit.script gpu version

script_cell_gpu = torch.jit.script(model_ft_gpu, (x_ft_gpu))
np.mean([timer(script_cell_gpu,x_ft.cuda()) for _ in range(100)])

• runtime

---

마치며..

이번 포스팅에서는 TorchScript와 Pytorch JIT에 대해 간단하게 알아보았다.

최근에는 Torchscript를 Just-In-Time (JIT) Compiler가 아닌, NVIDIA에서 개발한 TensorRT Compiler (Ahead-of-Time)를 이용하여 compile을 하는 추세이다. 혹은, pytorch model을 TorchScript가 아닌 ONNX format으로 변환한 후, 이를 TensorRT 등의 compiler를 통해 최적화하기도 한다.

---

reference

• PyTorch JIT and TorchScript
• TorchScript and PyTorch JIT / Deep Dive
• How to Convert a Model from PyTorch to TensorRT and Speed Up Inference
• Torch-TensorRT
• NVIDIA TensorRT – Inference 최적화 및 가속화를 위한 NVIDIA의 Toolkit
• pytorch 모델 저장과 ONNX 사용
• Torch-TensorRT
• LOADING A TORCHSCRIPT MODEL IN C++