import onnx
import torch
from onnx2keras import onnx_to_keras
# Load ONNX model
model = onnx.load('output/pytorch_onnx.onnx')
# Call the converter (input - is the main model input name, can be different for your model)
k_model = onnx_to_keras(model, ['input'])
Output -
W0827 20:22:03.981888 140701784454976 pooling_layers.py:33] Unable to use `same` padding. Add ZeroPadding2D layer to fix shapes.
W0827 20:22:03.985505 140701784454976 deprecation_wrapper.py:119] From /home/pat-011/.local/lib/python3.6/site-packages/keras/backend/tensorflow_backend.py:3661: The name tf.nn.max_pool is deprecated. Please use tf.nn.max_pool2d instead.
W0827 20:22:04.031678 140701784454976 pooling_layers.py:33] Unable to use `same` padding. Add ZeroPadding2D layer to fix shapes.
Model definition -
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 20, 5, 1)
self.conv2 = nn.Conv2d(20, 50, 5, 1)
self.fc1 = nn.Linear(4*4*50, 500)
self.fc2 = nn.Linear(500, 10)
def forward(self, x):
x = F.relu(self.conv1(x))
x = F.max_pool2d(x, 2, 2)
x = F.relu(self.conv2(x))
x = F.max_pool2d(x, 2, 2)
x = x.view(-1, 4*4*50)
x = F.relu(self.fc1(x))
x = self.fc2(x)
return F.softmax(x, dim=1)
model = Net().to(device)
Output -
Model definition -