torch.export: 현대적인 내보내기 시스템

TorchScript를 대체하는 시스템

torch.export는 PyTorch의 현대적인 모델 내보내기 시스템이야. 모델을 깔끔하고 표준화된 그래프로 포착해 직렬화하며, 배포·양자화 또는 ONNX, ExecuTorch, Core ML 같은 다른 실행 환경으로 변환하는 기반을 제공해. 새 프로젝트에서는 기존 TorchScript의 torch.jit.trace와 torch.jit.script를 대신해.

새 내보내기 시스템이 필요한 이유

TorchScript는 사용하기 어려웠어. jit.trace는 Python 제어 흐름을 보지 못하고, jit.script는 제한된 Python 부분집합을 요구해. 오류 메시지도 이해하기 힘들기로 유명했지. torch.export는 torch.compile과 같은 Dynamo 기반 그래프 포착을 사용해 더 넓은 코드를 지원하고 오류도 훨씬 명확하게 알려 줘.

계약

torch.export.export에 모델과 예시 입력을 넘겨. 이 입력으로 모델을 실행하며 그래프를 포착한 뒤 ExportedProgram을 반환해. 내보낸 그래프는 온전히 직렬화할 수 있어서 디스크에 저장하고 원본 Python 클래스 정의 없이 다시 불러올 수 있어.

torch.export로 얻는 것

실행 환경 이식성: 같은 내보낸 프로그램을 ONNX, ExecuTorch(모바일), Core ML, TensorRT로 변환할 수 있어.
양자화: torchao의 현대적인 양자화 기법을 내보낸 프로그램에 적용할 수 있어.
최적화: 상수 접기와 죽은 코드 제거 같은 그래프 패스를 내보낸 표현에 적용할 수 있어.
버전 관리: 내보내기 형식은 PyTorch 버전 사이에서 안정적으로 유지되도록 설계됐어.

TorchScript가 드물게 필요한 경우

일부 기존 임베디드 경로처럼 torch.export 실행 환경을 아직 사용할 수 없는 곳에 배포할 때는 TorchScript가 필요할 수 있어. 하지만 새 프로젝트의 기본 선택은 torch.export야.

Code

모델 내보내기와 불러오기·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)

# Export
exported = torch.export.export(model, (example,))

# Save — no Python class definition needed on load
torch.export.save(exported, "/tmp/tiny.pt2")

# Load (in another process / machine)
loaded = torch.export.load("/tmp/tiny.pt2")
y = loaded.module()(example)        # call .module() to get a callable
print(y.shape)                       # torch.Size([1, 4])

동적 배치 크기의 내보내기·python

import torch
from torch.export import Dim

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)

# Tell torch.export that dim 0 (batch) is dynamic
batch = Dim("batch", min=1, max=128)
exported = torch.export.export(
    model, (example,), dynamic_shapes=({0: batch},),
)

# Now the same exported program can run on different batch sizes
loaded = torch.export.load
torch.export.save(exported, "/tmp/dyn.pt2")
loaded = torch.export.load("/tmp/dyn.pt2")
print(loaded.module()(torch.randn(32, 10)).shape)   # torch.Size([32, 4])

ONNX 내보내기: 현대적인 dynamo 경로·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 (the modern path)
torch.onnx.export(
    model,
    (example,),
    "/tmp/tiny.onnx",
    input_names=["x"],
    output_names=["y"],
    dynamic_axes={"x": {0: "batch"}, "y": {0: "batch"}},
)

TorchScript: 구식 참조·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)

# Method 1: tracing — records ops from this example
traced = torch.jit.trace(model, example)
traced.save("/tmp/tiny_traced.pt")

# Method 2: scripting — compiles Python (restricted subset)
scripted = torch.jit.script(model)
scripted.save("/tmp/tiny_scripted.pt")

# Both can be loaded without Python — but for new projects, prefer torch.export

Exercise

같은 TinyMLP를 torch.export, torch.jit.trace, torch.jit.script 세 방식으로 내보내고 모두 저장해 봐. 각 파일을 불러와 추론 한 번의 시간과 파일 크기를 비교해. 세 방식이 모두 작동하는지 확인하고, 더 까다로운 입력에서 하나가 실패한다면 오류 메시지가 원인을 얼마나 잘 설명하는지도 비교해.