직렬화: state_dict에서 .pt2까지

네 가지 직렬화 형식과 쓰임새

state_dict(.pth / .pt): 매개변수와 버퍼를 담은 Python OrderedDict야. 가장 유연하지만 불러올 때 모델 클래스가 필요해. 학습과 연구의 기본 선택이야.
torch.export(.pt2): 계산 그래프와 가중치를 함께 내보내. 원본 Python 클래스 없이 불러올 수 있는 현대적인 배포 경로야.
TorchScript(.pt): 기존의 그래프와 가중치 형식이야. 여전히 널리 지원되지만 새 코드에서는 구식 경로로 봐.
ONNX(.onnx): 프레임워크 공통 표준이야. ONNX Runtime, TensorRT, OpenVINO, 브라우저의 onnxruntime-web처럼 PyTorch가 아닌 실행 환경에 배포할 때 적합해.

각 형식의 절충점

state_dict는 PyTorch 버전이 달라도 옮겨 쓰기 쉽고 모델 리팩터링에도 비교적 잘 견디지만 모델을 만들 Python 클래스가 필요해. torch.export와 TorchScript는 원본 클래스 없이 실행할 수 있지만 PyTorch 계열 실행 환경에 묶여. ONNX는 여러 실행 환경으로 옮기기 쉽지만 PyTorch에 특화된 최적화를 일부 잃을 수 있어.

결정 트리

추가 학습을 위해 공유할 때 → state_dict.
PyTorch 실행 환경에 배포할 때 → torch.export(.pt2).
ONNX Runtime, TensorRT, 브라우저에 배포할 때 → Dynamo 경로를 통한 ONNX.
모바일에 배포할 때 → 뒤에서 다룰 ExecuTorch.
Apple 플랫폼에 배포할 때 → 뒤에서 다룰 Core ML.

Code

state_dict: 유연하지만 Python 클래스가 필요해·python

import torch
import torch.nn as nn

class TinyMLP(nn.Module):
    def __init__(self): super().__init__(); self.fc = nn.Linear(10, 4)
    def forward(self, x): return self.fc(x)

model = TinyMLP()
torch.save(model.state_dict(), '/tmp/tiny.pth')

# To load — must have the class definition
m2 = TinyMLP()
m2.load_state_dict(torch.load('/tmp/tiny.pth', weights_only=True))
m2.eval()

torch.export (.pt2): 원본 클래스 없이 쓰는 현대적 형식·python

import torch
import torch.nn as nn

class TinyMLP(nn.Module):
    def __init__(self): super().__init__(); self.fc = nn.Linear(10, 4)
    def forward(self, x): return self.fc(x)

model = TinyMLP().eval()
example = torch.randn(1, 10)

exported = torch.export.export(model, (example,))
torch.export.save(exported, '/tmp/tiny.pt2')

# Load WITHOUT the class definition
loaded = torch.export.load('/tmp/tiny.pt2')
y = loaded.module()(example)        # call .module() to get a callable
print(y.shape)                       # torch.Size([1, 4])

ONNX 내보내기: PyTorch 밖의 실행 환경용·python

import torch

class TinyMLP(torch.nn.Module):
    def __init__(self): super().__init__(); self.fc = torch.nn.Linear(10, 4)
    def forward(self, x): return self.fc(x)

model = TinyMLP().eval()
example = torch.randn(1, 10)

# In PyTorch 2.x, torch.onnx.export defaults to dynamo=True (modern path)
torch.onnx.export(
    model, (example,), '/tmp/tiny.onnx',
    input_names=['x'], output_names=['y'],
    dynamic_axes={'x': {0: 'batch'}, 'y': {0: 'batch'}},
)

# Then run with ONNX Runtime (separately installed)
# pip install onnxruntime
import onnxruntime as ort
sess = ort.InferenceSession('/tmp/tiny.onnx')
out = sess.run(None, {'x': example.numpy()})
print(out[0].shape)                  # (1, 4)

Exercise

같은 TinyMLP를 state_dict, torch.export, ONNX 세 형식으로 직렬화해 봐. 각각 다시 불러올 수 있고 같은 입력에서 fp32 허용 오차 안의 같은 출력을 내는지 검증해. 파일 크기도 비교해. 작은 모델에서는 세 파일의 크기가 대체로 비슷할 거야.