Skip to content

Instantly share code, notes, and snippets.

@tkeyo

tkeyo/onnx_export.py Secret

Created July 6, 2021 21:26
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save tkeyo/021cd1aebde41f52917b2acc64b2ea59 to your computer and use it in GitHub Desktop.
Save tkeyo/021cd1aebde41f52917b2acc64b2ea59 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
onnx_export.py
import torch
torch.onnx.export(
final_model,
torch.randn(1, 3, 256, 256),
'models/hot_dog_model_resnet18_256_256.onnx',
do_constant_folding=True,
export_params=True,
input_names=['image_1_3_256_256'],
output_names=['hot_dog'],
opset_version=11
)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment