torch.export — modern Export System

TorchScript 의 대체

torch.export 가 PyTorch 의 modern model export system. 너 model 을 깔끔, 표준화된 graph representation 으로 capture, serialize, deploy, quantize, 또는 다른 backend (ONNX, ExecuTorch, CoreML) 로 lowering 가능. 새 project 에 옛 TorchScript (torch.jit.trace / torch.jit.script) 대체.

왜 새 export system

TorchScript 가 어려웠음. jit.trace 가 Python control flow 못 봄. jit.script 가 제한된 Python subset 요구. 에러 메시지가 악명 높게 cryptic. torch.export 가 torch.compile 와 같은 dynamo-based graph capture 사용, 훨씬 더 나은 coverage 와 훨씬 깔끔한 에러.

계약

torch.export.export 에 model 과 example input 줘. 그 input 으로 model 돌리고 결과 graph capture, ExportedProgram 반환. exported graph 가 fully serializable — disk 에 저장하고 원본 Python 클래스 정의 없이 로드.

torch.export 가 사주는 것

backend 이식성 — 같은 exported program 이 ONNX, ExecuTorch (mobile), CoreML, TensorRT 로 lowering.
quantization — torchao 의 modern quant 기법이 exported program 에 동작.
최적화 — graph pass (constant folding, dead code elimination) 가 exported representation 에 동작.
versioning — export format 이 PyTorch version 들에 안정 의도.

TorchScript — (드물게) 여전히 잡을 때

torch.export 의 runtime 이 아직 사용 가능 안 한 환경에 deploy 위 (일부 legacy embedded path). 새 project 에 default torch.export.

Code

model export 와 load·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])

dynamic batch size 의 export·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 export — modern dynamo path·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 — legacy 참조·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 를 세 방법으로 export: torch.export, torch.jit.trace, torch.jit.script. 다 저장. 각 로딩 + 한 inference call 시간. 파일 size diff. torch.export 와 TorchScript variant 모두 작동; 어느 거 ship 할지 까다로운 input 에 어떤 거 fail 하면 에러 메시지 읽어서 결정.