用 TorchScript 匯出模型

Scalable AI Models with PyTorch Lightning

Sergiy Tkachuk

Director, GenAI Productivity

什麼是 TorchScript？

不依賴 Python
適合生產環境，高效
範例：
- 佈署於行動裝置

什麼是 TorchScript？

不依賴 Python
生產環境高效
範例：
- 佈署於行動裝置
- 生產系統中的高效推論

手機與筆電，齒輪代表生產環境

將模型轉換為 TorchScript

兩種轉換方式：

torch.jit.trace：用範例輸入追蹤執行
torch.jit.script：分析 Python 程式碼進行編譯

何時使用：

簡單模型用 trace
有控制流程（如迴圈）的模型用 script

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def forward(self, x):
        return x * 2

model = SimpleModel()
scripted_model = torch.jit.script(model)

儲存與載入 TorchScript 模型

儲存模型：
- torch.jit.save：將 scripted 模型存成檔案
載入模型：
- torch.jit.load：載回模型以進行推論

# Save the model
torch.jit.save(scripted_mod,"model.pt")

# Load the model
loaded_model=torch.jit.load("model.pt")

使用 TorchScript 進行推論

步驟：
- 載入 TorchScript 模型
- 將輸入傳入模型取得預測
- 輸出與 PyTorch 預測相同

範例輸入：

輸入張量：[1.0, 2.0, 3.0]

範例輸出：

輸出張量：[2.0, 4.0, 6.0]

# Perform inference
input_arr = [1.0, 2.0, 3.0]
input_tensor = torch.tensor(input_arr)

output = loaded_model(input_tensor)
print(output)

TorchScript 重點摘要

torch.jit.trace：適用靜態模型
torch.jit.script：可處理動態控制流程
torch.jit.save：儲存 scripted 模型
torch.jit.load：載入以進行推論

一起來練習吧！

Scalable AI Models with PyTorch Lightning