convert_to_onnx.py

import torch
import torchvision.models as models
import torch.onnx as onnx

# 載入 ResNet-18 模型
model = models.resnet50(pretrained=True)

# 將模型設定為評估模式
model.eval()

# 定義輸入張量
x = torch.randn(1, 3, 224, 224)  # 假設輸入圖片尺寸為 224x224，通道數為 3

# 將模型轉換為 ONNX 格式
onnx_file = "resnet50.onnx"
onnx.export(model, x, onnx_file, input_names=['input'])

print("ResNet-18 模型已轉換為 ONNX 格式，並保存為", onnx_file)