ExecuTorch와 Arm을 활용한 엣지 기기용 PyTorch 모델 배포 가이드

import torch
import torch.nn as nn

class TinyRPS(nn.Module):
    def __init__(self):
        super().__init__()
        self.body = nn.Sequential(
            nn.Conv2d(1, 16, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(16, 32, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2),
        )
        self.head = nn.Sequential(
            nn.Flatten(),
            nn.Linear(64 * 7 * 7, 128),
            nn.ReLU(inplace=True),
            nn.Linear(128, 3),
        )

    def forward(self, x):
        x = self.body(x)
        x = self.head(x)
        return x

from executorch import exir
from torch.export import export

def export_to_pte(model: nn.Module, out_path: str, img_size: int) -> None:
    model.eval()
    example = torch.zeros(
        1, 1, img_size, img_size, dtype=torch.float32
    )
    # Export with PyTorch’s exporter
    exported = export(model, (example,))
    edge = exir.to_edge(exported)
    prog = edge.to_executorch()
    with open(out_path, "wb") as f:
        f.write(prog.buffer)
    print(f"[export] wrote {out_path}")