NonImplementedError on using torch.onnx.export

Question

I am trying to convert a pre-saved PyTorch model into a TensorFlow one via ONNX. For now, the following code is to export the model into .onnx format. The neural network has 2 inputs, one hidden layer with 5 neurons and a scalar output.

Here's the code I'm working with:

import torch.nn as nn
from torch.autograd import Variable
import numpy as np

class Model(nn.Module):
    def __init__(self, n_h_layers, n_h_neurons, dim_in, dim_out, in_bound, out_bound):
        super(Model,self).__init__()
        self.n_h_layers=n_h_layers
        self.n_h_neurons=n_h_neurons
        self.dim_in=dim_in
        self.dim_out=dim_out
        self.in_bound=in_bound
        self.out_bound=out_bound
        layer_input = [nn.Linear(dim_in, n_h_neurons, bias=True)]
        layer_output = [nn.ReLU(), nn.Linear(n_h_neurons, dim_out, bias=True), nn.Hardtanh(in_bound, out_bound)]

        # hidden layer
        module_hidden = [[nn.ReLU(), nn.Linear(n_h_neurons, n_h_neurons, bias=True)] for _ in range(n_h_layers - 1)]
        layer_hidden = list(np.array(module_hidden).flatten())

        # nn model
        layers = layer_input + layer_hidden + layer_output
        self.model = nn.Sequential(*layers)

        print(self.model)
        

trained_nn=torch.load('path')            

trained_model=Model(1,5,2,1,-1,1)  
trained_model.load_state_dict(trained_nn,strict=False)

dummy_input=Variable(torch.randn(1,2))
torch.onnx.export(trained_model,dummy_input, 'file.onnx', verbose=True) 

I have two problems:

1. Running this snippet raises "NonImplementedError" in _forward_unimplemented in module.py as follows:

File ".../anaconda3/lib/python3.9/site-packages/torch/nn/modules/module.py", line 201, in _forward_unimplemented
    raise NotImplementedError

NotImplementedError

I am not aware with Exception handling in python and I do not know what I must change in order to tackle the error.

2. When I print trained_nn, this is what it gives me:

OrderedDict([('0.weight',
              tensor([[ 0.2035, -0.7679],
                      [ 1.6368, -0.4135],
                      [-0.0908, -0.2335],
                      [ 1.3731, -0.3135],
                      [ 0.6361,  0.2521]])),
             ('0.bias', tensor([-1.6907,  0.7262,  1.4032,  1.2551,  0.8013])),
             ('2.weight',
              tensor([[-0.4603, -0.0719,  0.4082, -1.0235, -0.0538]])),
             ('2.bias', tensor([-1.1568]))])

However, printing trained_model.state_dict() gives me a neural network with a completely different set of weights and biases, although I believe that it should be giving me the exact same model as before as this is what I need to save as onnx file?

OrderedDict([('model.0.weight',
              tensor([[ 0.4817,  0.0928],
                      [-0.4313,  0.1253],
                      [ 0.6681, -0.4029],
                      [ 0.6474,  0.0029],
                      [-0.4663,  0.5029]])),
             ('model.0.bias',
              tensor([-0.2292,  0.6674, -0.3755,  0.0778,  0.0527])),
             ('model.2.weight',
              tensor([[-0.2097, -0.3029,  0.2792,  0.2596,  0.1362]])),
             ('model.2.bias', tensor([-0.1835]))])

Not sure what mistakes I'm making. Any help is appreciated.

Answer 1

1. When you are making a subclass of nn.Module you need to implement forward method. In your case you need to add:

class Model(nn.Module):
    def __init__(self, n_h_layers, n_h_neurons, dim_in, dim_out, in_bound, out_bound):
        super(Model, self).__init__()

    ...

    def forward(self, x):
        return self.model(x)

2. The names of parameters does not match:

model.0.weight != 0.weight

model.0.bias != 0.bias

prefix model is missed.

So when you call load_state_dict() with strict=False the parameters will not be used.

You can rename the parameters to match the model:

trained_nn = torch.load('path')            
trained_nn = {f'model.{name}': w for name, w in trained_nn.items()}
trained_model.load_state_dict(trained_nn, strict=True)