Skip to content

Instantly share code, notes, and snippets.

@justinchuby

justinchuby/export_pytorch_onnx.py

Last active January 16, 2024 17:57
Show Gist options
  • Save justinchuby/0cf5ec03a0beb1ce3cb6aa5cb24bf06f to your computer and use it in GitHub Desktop.
Save justinchuby/0cf5ec03a0beb1ce3cb6aa5cb24bf06f to your computer and use it in GitHub Desktop.
Download ZIP
Test PyTorch ONNX export
Raw
export_pytorch_onnx.py
import torch
import onnx
class Model(torch.nn.Module):
def forward(self, x):
return torch.nn.functional.hardsigmoid(x)
exported = torch.onnx.dynamo_export(Model(), torch.randn(1, 3, 224, 224))
print(onnx.printer.to_text(exported.model_proto))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment