[Pytorch] TorchScript

Pytorch 기본 모델 형식은 태생적인 단점을 많이 가지고 있는 편이다.
파이썬 코드로 실행되고, 사용법이 간단하고, 학습이나 배포도 간단한 것이 장점이나. 최적의 성능을 내지 못한다는게 고질적인 단점이 있었다.

그래서 나온 것이 Torchscript라는 모델 파일 형식이다.

torch에서 제공하는 torchscript 기능 지원을 활용하면, 컴파일 과정을 거쳐서 python&pytorch 종속성을 제거하고, 최소한의 의미만을 담는 script를 만들 수 있다.
다만, 이 자체로는 완성된 바이너리가 아님에 유의한다. 바이트코드처럼 중간 표현 언어로 변환하는 것이다.

실제로 torchscript를 돌리려면 이걸 다시 torch로 로드해서 실행하거나, TensorRT 같은 진짜 native 바이너리 형태로 컴파일해서 실행해야 한다.

예제 샘플

먼저 예제 모델을 하나 준비해보겠다.
그냥 길이 2짜리 입력을 받아서 길이 3짜리를 반환하는 간단한 모델이다.

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(2, 3)

    def forward(self, x):
        return self.linear(x)

from model import SimpleModel
import torch
import torch.nn as nn
import torch.optim as optim

model = SimpleModel()

print(model)

x = torch.tensor([[1.0, 1.0], [2.0, 2.0], [3.0, 3.0], [4.0, 4.0]])
y = torch.tensor([[3.0, 5.0, 7.0], [5.0, 7.0, 9.0], [7.0, 9.0, 11.0], [9.0, 11.0, 13.0]])

criterion = nn.MSELoss()  # 평균제곱오차
optimizer = optim.SGD(model.parameters(), lr=0.01)  # 확률적 경사하강법

# 10번 반복해서 학습 (epoch)
for epoch in range(1000):
    optimizer.zero_grad()         # 기울기 초기화
    outputs = model(x)            # 모델에 입력값 전달
    loss = criterion(outputs, y)  # 손실 계산
    loss.backward()               # 역전파
    optimizer.step()              # 파라미터 업데이트

    if (epoch+1) % 100 == 0:
        print(f'Epoch [{epoch+1}/1000], Loss: {loss.item():.4f}')

# 결과 확인
print("학습된 가중치:", model.linear.weight.tolist())
print("학습된 편향:", model.linear.bias.tolist())

# 테스트 
print("기대한 출력", y.tolist())
print("실제의 출력",  model(torch.tensor([[1.0, 1.0]])).tolist())

# 3가지 방식으로 모델 저장
torch.save(model.state_dict(), "state_dict.pth")
torch.save(model, "model_full.pth")
checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
    # ... 기타 필요한 정보
}
torch.save(checkpoint, "checkpoint.pth")

그래서 1.0, 1.0을 넣으면 3,5,7에 근사하는 값이 나오도록 세팅해뒀다.

이걸 torchscript로 컴파일해보자.

2가지 컴파일 방법: trace vs script

torchscript에는 2가지 컴파일 방법이 있다.
근데 이게 방법만 다른게 아니라 실제 결과물의 수준도 다를 수 있기 때문에 그 장단점과 특징을 잘 알고 선택해야 한다. 딱 뭐를 써야한다는 일반적인 지침은 존재하지 않는 것 같다.

일단 script는 가장 정직하고 일반적인 형태의 컴파일 방법이다. 그냥 python 코드를 읽어서, 정적으로 컴파일하는 전형적인 방법론을 취한다.

반면에 trace는 모델을 실제로 실행하면서 그 실행 흐름에 맞게 최적화된 코드를 만들어내는 방식을 취한다. 진짜 JIT 컴파일에 가까운 방법이라고 할 수 있겠다.

trace 컴파일

trace 컴파일은 그렇게 어렵진 않다.
실제 모델 실행에 필요한 입력값을 대충 하나 만들고, 밀어넣어서 trace 컴파일을 돌리면 된다.

import trace
from model import SimpleModel
import torch

model = SimpleModel()

model.load_state_dict(torch.load("state_dict.pth"))

print(model)

# 추론 모드 전환
model.eval()

# Tracing 방식 - 예제 입력으로 모델 실행 경로 추적
example_input = torch.tensor([[1.0, 2.0]])
traced_model = torch.jit.trace(model, example_input)

traced_model.save("traced_model.pt")

그러고 만들어지는 반환값 traced_model이 trace 모드로 컴파일된 모델 데이터다.
저걸 저장해서 사용하면 된다.

실행할 때는 jit.load를 통해서 로드해서 사용할 수 있다.

from model import SimpleModel
import torch

loaded_model = torch.jit.load("traced_model.pt")

print(loaded_model)

print(loaded_model(torch.tensor([[1.0, 1.0]])).tolist())

경우에 따라 값이 약간 왜곡될 수도 있긴 한데, 이 경우에는 동일한 값이 나왔다.

script 컴파일

script 컴파일의 경우에는 비교적 더 간단하다.
실제로 모델을 돌릴 필요도 없기 때문이다.

from model import SimpleModel
import torch

loaded_model = torch.jit.load("scripted_model.pt")

print(loaded_model)

print(loaded_model(torch.tensor([[1.0, 1.0]])).tolist())

이 경우에도 일단은 동등한 결과가 나왔다.

trace vs script

그래서, 2가지 방법 중 무엇을 선택해야할까?
사실 이것도 경우에 따라 다르다는 말을 할 수밖에 없을 것 같다.

trace는 해당 계산 흐름에 맞춰서 최적화된 형태로 정적 빌드를 수행한다. 해당 예제 실행에 사용되지 않은 불필요한 그래프를 전부 제거한다.
그래서 실행 흐름이 한가지라면 무엇보다 빠르고 뛰어난 결과를 도출할 수 있다.
하지만 다양한 형태의 실행 흐름을 제공하게끔 만들기 어렵고, 입력 크기 또한 한가지로 제한된다는 단점이 있다.

script는 코드를 정적으로 분석해서 최적화된 형태로 컴파일을 시도한다. 이건 특히, 모델 코드에 복잡한 조건문이나 반복문 같은 것이 있는 패턴에 대해서는 더욱 유효하게 동작한다.
예를 들면, 이런 식의 변칙적인 코드 말이다.

import torch
import torch.nn as nn

class SampleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(2, 3)

    def forward(self, x):
        # 동적 제어 흐름 예시
        if x.sum() > 0:
            return self.linear(x)
        else:
            return torch.zeros_like(x)

하지만 script 코드의 정적 분석 수준은 낮다. 모든 python 신택스나 코드 패턴에 유효하게 동작하는 것이 아니기 때문이다. 최적화가 되는 한도 내에서 강박적으로 코드를 작성해야 한다는게 좀 치명적인 부분이다.
상속을 써도 안되고, lamdba를 써도 안되고, union도 안되고, 동적 타입도 안되고, 사실 뭐 되는게 없다.
코드 퀄리티와 script 컴파일 가능성을 동시에 가져가는게 거의 불가능하다.

trace with script

양쪽 모두의 장점을 가져가고 싶다면, trace와 script를 동시에 섞어서 사용하는 것도 가능하다.
예를 들면, 이런 식이다.

class MyRNNLoop(torch.nn.Module):
    def __init__(self):
        super(MyRNNLoop, self).__init__()
        self.cell = torch.jit.trace(MyCell(scripted_gate), (x, h))

    def forward(self, xs):
        h, y = torch.zeros(3, 4), torch.zeros(3, 4)
        for i in range(xs.size(0)):
            y, h = self.cell(xs[i], h)
        return y, h

rnn_loop = torch.jit.script(MyRNNLoop())
print(rnn_loop.code)

세부 모델 동작들에 대해서는 trace로 컴파일을 하고, 분기 제어에 대해서는 script 컴파일을 하도록 하는 것이다.
이러면 컴파일을 할때 script가 trace된 스크립트를 인라인해서 적절히 처리한다.

예제에 사용한 전체 코드베이스는 다음 링크에서 확인할 수 있다.
https://github.com/myyrakle/ml_examples/tree/master/pytorch/torchscript-basic

참조
https://tutorials.pytorch.kr/recipes/torchscript_inference.html
https://happy-jihye.github.io/dl/torch-2/
https://ppwwyyxx.com/blog/2022/TorchScript-Tracing-vs-Scripting/
https://docs.pytorch.org/docs/stable/jit.html
https://docs.pytorch.org/docs/stable/jit_language_reference_v2.html