Skip to content

Instantly share code, notes, and snippets.

@twmht
Created July 28, 2023 05:56
Show Gist options
  • Save twmht/1dc0c41a80478c2de6f215824a7b8e14 to your computer and use it in GitHub Desktop.
Save twmht/1dc0c41a80478c2de6f215824a7b8e14 to your computer and use it in GitHub Desktop.
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment