Build a Mock Model in PyTorch with a convolution and a reduceMean layer
import torch
import torch.onnx as torch_onnx
from torch.autograd import Variable
class MLP(torch.nn.Module):
def init(self):
super(MLP,self).init()
self.fc1 = torch.nn.Linear(1,1)
def forward(self,din):
dout = torch.nn.functional.relu(self.fc1(din))
return dout
Use this an input trace to serialize the model
model_onnx_path = “mlpsimple1.onnx”
model = MLP()
print (model.state_dict())
Export the model to an ONNX file
dummy_input = Variable(torch.randn(1,1))
output = torch.onnx.export(model,
dummy_input,
model_onnx_path,
export_params=True,
verbose=True)
print(“Export of torch_model.onnx complete!”)