I got always same result form RNN model exported to ONNX

Question

I exported pytorch RNN model to ONNX with torchscript but onnx return always same result. i got different results with pytorch, torchscript here is my code

class RNN(nn.Module):
    
    def __init__(self, input_size, hidden_size, num_classes, n_layers=2): 
        
        super(RNN, self).__init__()
        
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.num_classes = num_classes
        self.n_layers = n_layers
        self.hidden = None
        self.notes_encoder = nn.Linear(in_features=input_size, out_features=hidden_size)
        self.bn = nn.BatchNorm1d(hidden_size)
        self.lstm = nn.LSTM(hidden_size, hidden_size, n_layers)
        self.logits_fc = nn.Linear(hidden_size, num_classes)

    def forward(self, inMelody): 
        current_sequence_input = inMelody 
        #final_output_sequence =[]
        final_output_sequence =[current_sequence_input.data.squeeze(1)]
        
        for i in  range(200): 
            notes_encoded = self.notes_encoder(current_sequence_input)  
            notes_encoded_rolled = notes_encoded.permute(1,2,0).contiguous() 
            notes_encoded_norm_drop  = self.bn(notes_encoded_rolled) 
            notes_encoded_complete = notes_encoded_norm_drop.permute(2,0,1) 

            notes_encoded_sqz = notes_encoded_complete[-1:]  
            outputs, _= self.lstm(notes_encoded_sqz, self.hidden) 
            outputs_drop = self.bn(outputs.permute(1,2,0).contiguous())

            logits = self.logits_fc(outputs_drop.permute(2,0,1)) 
            logits = logits.transpose(0, 1).contiguous() 
            neg_logits = (1 - logits)
            
            binary_logits = torch.stack((logits, neg_logits), dim=3).contiguous()
            logits_flatten = binary_logits.view(-1, 2)
            output = logits_flatten    
            probabilities = nn.functional.softmax(output.div(0.8), dim=1) 

            current_sequence_input = torch.multinomial(probabilities, 1, replacement=True).squeeze().unsqueeze(0).unsqueeze(1) 
            current_sequence_input = current_sequence_input.float()  
            final_output_sequence.append(current_sequence_input.data.squeeze(1)) 

        outMelody = torch.cat(final_output_sequence, dim=0)
        return outMelody 

rnn = RNN(input_size=88, hidden_size=512, num_classes=88)
rnn = rnn.cuda()
rnn.load_state_dict(torch.load('music_model_padfront_regularized.pth')) 
rnn.eval()

sample = rnn(input_sequences_batch_var.cuda()).cpu().numpy().transpose() 
plt.imshow(sample)

export to ONNX Code

import torch.onnx

Mrnn = RNN(input_size=88, hidden_size=512, num_classes=88)
Mrnn = Mrnn.cuda()
Mrnn.load_state_dict(torch.load('music_model_padfront_regularized.pth')) 
Mrnn.eval()

my_modifier_scripted = torch.jit.script(Mrnn)
inMelody =  torch.empty(1, 1, 88, dtype = torch.float32)
example_input =(inMelody.cuda()) 


input_names = ['inMelody'] 
output_names = ['outMelody']

torch.onnx.export(my_modifier_scripted, 
                  example_input,
                  "MusicGenerator555.onnx",
                  input_names = input_names,   
                  output_names = output_names,
                  dynamic_axes= {'inMelody' : [0]})

here is some warning

UserWarning: Exporting a model to ONNX with a batch_size other than 1, with a variable length with LSTM can cause an error when running the ONNX model with a different batch size. Make sure to save the model with a batch size of 1, or define the initial states (h0/c0) as inputs of the model. 

UserWarning: The shape inference of prim::Constant type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function. (Triggered internally at ..\torch\csrc\jit\passes\onnx\shape_type_inference.cpp:1888.)

onnx test code

import os
import torch
import psutil
import onnxruntime as ort
import numpy as np

ort_session = ort.InferenceSession("MusicGenerator555.onnx", providers=['TensorrtExecutionProvider', 'CUDAExecutionProvider'])
sample = ort_session.run(['outMelody'], {'inMelody': input_sequences_batch_var.numpy()}) 

lastRes = np.squeeze(sample,  axis=0)                                         
plt.imshow(lastRes.transpose())

this onnx model returns alwasy same result :( plese let know newbie what is wrong~

Here is Colab link https://colab.research.google.com/drive/1z_rZu2IHWGfB4OY6NloSR8T2JQ6cIe_d?hl=ko#scrollTo=aRdzA2not-CF

Appreciate any kind of help! Thank you in advance.

I got always same result form RNN model exported to ONNX

0 Answers0