Cryptic error message when converting torch model to onnx

Harry-Garrison: <p>I’ve been trying to export a model from torch to onnx using the torch.jit.script(model()) command. I have also placed various print statements throughout my code just to make sure that it is running. However, at some point execution stops and I get the following error message:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "C:\Users\User\Speaker-Diarization\uisrnn_plain_functions_old.py", line 269, in &lt;module&gt;
    model = torch.jit.script(predict(x,model_args))
  File "C:\Users\User\Python\lib\site-packages\torch\jit\_script.py", line 901, in script
    qualified_name = _qualified_name(obj)
  File "C:\Users\User\Python\lib\site-packages\torch\_jit_internal.py", line 793, in _qualified_name
    raise RuntimeError("Could not get name of python class object")
RuntimeError: Could not get name of python class object
</code></pre>
<p>As of this point I have no idea what goes wrong. I tested it on two versions of Python: Pyhton 3.6 (Anaconda version) and Python 3.7 with PyTorch version 1.0 and 1.7 respectively. On the Anaconda version I got the following error:</p>
<pre><code class="lang-auto">
Traceback (most recent call last):
  File "C:\Users\User\Speaker-Diarization\uisrnn_plain_functions_old.py", line 269, in &lt;module&gt;
    model = torch.jit.script(predict(x,model_args))
  File "C:\Users\User\Anaconda3\lib\site-packages\torch\jit\__init__.py", line 689, in script
    ast = get_jit_ast(fn, is_method=False)
  File "C:\Users\User\Anaconda3\lib\site-packages\torch\jit\frontend.py", line 141, in get_jit_ast
    source = dedent(inspect.getsource(fn))
  File "C:\Users\User\Anaconda3\lib\inspect.py", line 973, in getsource
    lines, lnum = getsourcelines(object)
  File "C:\Users\User\Anaconda3\lib\inspect.py", line 955, in getsourcelines
    lines, lnum = findsource(object)
  File "C:\Users\User\Anaconda3\lib\inspect.py", line 768, in findsource
    file = getsourcefile(object)
  File "C:\Users\User\Anaconda3\lib\inspect.py", line 684, in getsourcefile
    filename = getfile(object)
  File "C:\Users\User\Anaconda3\lib\inspect.py", line 666, in getfile
    'function, traceback, frame, or code object'.format(object))
TypeError: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] is not a module, class, method, function, traceback, frame, or code object
</code></pre>
<p>Since there aren’t many resources on this, I tried to replicate the methods found on PyTorch’s website without success. My code is as follows:</p>
<pre><code class="lang-auto">import numpy as np
import torch
from torch import autograd
from torch import nn
from torch import optim
import torch.nn.functional as F
import uisrnn
from uisrnn import loss_func
from uisrnn import utils

import onnx
import onnxruntime


class CoreRNN(nn.Module):
  """The core Recurent Neural Network used by UIS-RNN."""

  def __init__(self, input_dim, hidden_size, depth, observation_dim, dropout=0):
    super(CoreRNN, self).__init__()
    self.hidden_size = hidden_size
    if depth &gt;= 2:
      self.gru = nn.GRU(input_dim, hidden_size, depth, dropout=dropout)
    else:
      self.gru = nn.GRU(input_dim, hidden_size, depth)
    self.linear_mean1 = nn.Linear(hidden_size, hidden_size)
    self.linear_mean2 = nn.Linear(hidden_size, observation_dim)

  def forward(self, input_seq, hidden=None):
    output_seq, hidden = self.gru(input_seq, hidden)
    if isinstance(output_seq, torch.nn.utils.rnn.PackedSequence):
      output_seq, _ = torch.nn.utils.rnn.pad_packed_sequence(output_seq, batch_first=False)
    mean = self.linear_mean2(F.relu(self.linear_mean1(output_seq)))

    #print("Mean and hidden from Core", mean.shape, hidden.shape)
    
    return mean, hidden




class BeamState:
  """Structure that contains necessary states for beam search."""

  def __init__(self, source=None):
    if not source:
      self.mean_set = []
      self.hidden_set = []
      self.neg_likelihood = 0
      self.trace = []
      self.block_counts = []
      
    else:
      self.mean_set = source.mean_set.copy()
      self.hidden_set = source.hidden_set.copy()
      self.trace = source.trace.copy()
      self.block_counts = source.block_counts.copy()
      self.neg_likelihood = source.neg_likelihood

  def append(self, mean, hidden, cluster):
    #print("Mean and hidden from Beamstate", mean.shape, hidden.shape, cluster)
    """Append new item to the BeamState."""
    self.mean_set.append(mean.clone()) #old code
    self.hidden_set.append(hidden.clone()) #old code
    #self.mean_set.append(mean) #old code
    #self.hidden_set.append(mean) #old code
    self.block_counts.append(1)
    self.trace.append(cluster)

def uisrnn_onnx(model, ort_session, init_input, rnn_init_hidden):

    ort_inputs = {ort_session.get_inputs()[0].name: init_input}
    mean, hidden = ort_session.run(None, ort_inputs)
    return mean, hidden


#@torch.jit.script    
def _update_beam_state(beam_state, look_ahead_seq, cluster_seq):

    """Update a beam state given a look ahead sequence and known cluster
    assignments.

    Args:
      beam_state: A BeamState object.
      look_ahead_seq: Look ahead sequence, size: look_ahead*D.
        look_ahead: number of step to look ahead in the beam search.
        D: observation dimension
      cluster_seq: Cluster assignment sequence for look_ahead_seq.

    Returns:
      new_beam_state: An updated BeamState object.
    """

    loss = 0
    new_beam_state = BeamState(beam_state)
    for sub_idx, cluster in enumerate(cluster_seq):
      if cluster &gt; len(new_beam_state.mean_set):  # invalid trace
        new_beam_state.neg_likelihood = float('inf')
        break
      elif cluster &lt; len(new_beam_state.mean_set):  # existing cluster
        last_cluster = new_beam_state.trace[-1]
        loss = loss_func.weighted_mse_loss(input_tensor=torch.squeeze(new_beam_state.mean_set[cluster]), target_tensor=look_ahead_seq[sub_idx, :], weight=1 / (2 * sigma2)).cpu().detach().numpy() #old code
        #loss = loss_func.weighted_mse_loss(input_tensor=torch.squeeze(torch.tensor(new_beam_state.mean_set[cluster])), target_tensor=look_ahead_seq[sub_idx, :], weight=1 / (2 * sigma2)).cpu().detach().numpy()
        if cluster == last_cluster:
          loss -= np.log(1 - transition_bias)
        else:
          loss -= np.log(transition_bias) + np.log( new_beam_state.block_counts[cluster]) - np.log( sum(new_beam_state.block_counts) + crp_alpha)
        # update new mean and new hidden
        mean, hidden = rnn_model(look_ahead_seq[sub_idx, :].unsqueeze(0).unsqueeze(0), new_beam_state.hidden_set[cluster]) # old code
        #mean, hidden = uisrnn_onnx(onnx_model, ort_session, look_ahead_seq[sub_idx, :].unsqueeze(0).unsqueeze(0), new_beam_state.hidden_set[cluster]) #old code
        #mean, hidden = uisrnn_onnx(onnx_model, ort_session, look_ahead_seq[sub_idx, :].unsqueeze(0).unsqueeze(0).cpu().detach().numpy(), new_beam_state.hidden_set[cluster])
        new_beam_state.mean_set[cluster] = (new_beam_state.mean_set[cluster]*((np.array(new_beam_state.trace) == cluster).sum() -1).astype(float) + mean.clone()) / (np.array(new_beam_state.trace) == cluster).sum().astype(float)  #old code
        #new_beam_state.mean_set[cluster] = (new_beam_state.mean_set[cluster]*((np.array(new_beam_state.trace) == cluster).sum() -1).astype(float) + np.copy(mean)) / (np.array(new_beam_state.trace) == cluster).sum().astype(float)  # use mean to predict
        new_beam_state.hidden_set[cluster] = hidden.clone() #old code
        #new_beam_state.hidden_set[cluster] = hidden
        if cluster != last_cluster:
          new_beam_state.block_counts[cluster] += 1
        
        new_beam_state.trace.append(cluster)
      else:  # new cluster
        init_input = autograd.Variable(torch.zeros(observation_dim)).unsqueeze(0).unsqueeze(0).to(device) #old code
        #init_input = autograd.Variable(torch.zeros(observation_dim)).unsqueeze(0).unsqueeze(0).to(device).cpu().detach().numpy()
        #print(init_input.shape, "Shape of init_input")
        mean, hidden = rnn_model(init_input, rnn_init_hidden) #old code
        #mean, hidden = uisrnn_onnx(onnx_model, ort_session, init_input, rnn_init_hidden)
        loss = loss_func.weighted_mse_loss(input_tensor=torch.squeeze(mean), target_tensor=look_ahead_seq[sub_idx, :], weight=1 / (2 * sigma2)).cpu().detach().numpy() #old code
        #loss = loss_func.weighted_mse_loss(input_tensor=torch.squeeze(torch.tensor(mean)), target_tensor=look_ahead_seq[sub_idx, :], weight=1 / (2 * sigma2)).cpu().detach().numpy() 
        loss -= np.log(transition_bias) + np.log(crp_alpha) - np.log(sum(new_beam_state.block_counts) + crp_alpha)
        # update new min and new hidden
        mean, hidden = rnn_model(look_ahead_seq[sub_idx, :].unsqueeze(0).unsqueeze(0), hidden) #old code
        #mean, hidden = uisrnn_onnx(onnx_model, ort_session, look_ahead_seq[sub_idx, :].unsqueeze(0).unsqueeze(0), hidden) #old code
        #mean, hidden = uisrnn_onnx(onnx_model, ort_session, look_ahead_seq[sub_idx, :].unsqueeze(0).unsqueeze(0).cpu().detach().numpy(), hidden)
        new_beam_state.append(mean, hidden, cluster)
        
        
      new_beam_state.neg_likelihood += loss
    return new_beam_state


def _calculate_score(beam_state, look_ahead_seq):
    #print('Beamstate from calculate_score', beam_state)
    """Calculate negative log likelihoods for all possible state allocations
       of a look ahead sequence, according to the current beam state.

    Args:
      beam_state: A BeamState object.
      look_ahead_seq: Look ahead sequence, size: look_ahead*D.
        look_ahead: number of step to look ahead in the beam search.
        D: observation dimension

    Returns:
      beam_score_set: a set of scores for each possible state allocation.
    """

    look_ahead, _ = look_ahead_seq.shape
    beam_num_clusters = len(beam_state.mean_set)
    beam_score_set = float('inf') * np.ones(beam_num_clusters + 1 + np.arange(look_ahead))
    for cluster_seq, _ in np.ndenumerate(beam_score_set):
      updated_beam_state = _update_beam_state(beam_state, look_ahead_seq, cluster_seq)
      beam_score_set[cluster_seq] = updated_beam_state.neg_likelihood
    return beam_score_set

def load(filepath):
    """Load the model from a file.

    Args:
      filepath: the path of the file.
    """
    var_dict = torch.load(filepath)
    #rnn_model.load_state_dict(var_dict['rnn_state_dict'])
    rnn_init_hidden = nn.Parameter(torch.from_numpy(var_dict['rnn_init_hidden']).to(device))
    transition_bias = float(var_dict['transition_bias'])
    transition_bias_denominator = float(var_dict['transition_bias_denominator'])
    crp_alpha = float(var_dict['crp_alpha'])
    sigma2 = nn.Parameter(torch.from_numpy(var_dict['sigma2']).to(device))

    return var_dict


def predict_single(test_sequence, args):
    # check type
    if (not isinstance(test_sequence, np.ndarray) or
        test_sequence.dtype != float):
      raise TypeError('test_sequence should be a numpy array of float type.')
    # check dimension
    if test_sequence.ndim != 2:
      raise ValueError('test_sequence must be 2-dim array.')
    # check size
    test_sequence_length, observation_dim = test_sequence.shape
    if observation_dim != observation_dim:
      raise ValueError('test_sequence does not match the dimension specified '
                       'by args.observation_dim.')


    
    test_sequence = np.tile(test_sequence, (2, 1)) #args.test_iteration =2
    test_sequence = autograd.Variable(torch.from_numpy(test_sequence).float()).to(device)
    # bookkeeping for beam search
    beam_set = [BeamState()]
    for num_iter in np.arange(0, 2 * test_sequence_length, 1): #args.test_iteration= 2, args.look_ahead=1
      max_clusters = max([len(beam_state.mean_set) for beam_state in beam_set])
      
      look_ahead_seq = test_sequence[num_iter:  num_iter + 1, :] #args.look_ahead=1
      look_ahead_seq_length = look_ahead_seq.shape[0]
      score_set = float('inf') * np.ones(np.append(10, max_clusters + 1 + np.arange(look_ahead_seq_length))) #args.beam_size = 10
      for beam_rank, beam_state in enumerate(beam_set):
        beam_score_set = _calculate_score(beam_state, look_ahead_seq)
        score_set[beam_rank, :] = np.pad(beam_score_set, np.tile([[0, max_clusters - len(beam_state.mean_set)]], (look_ahead_seq_length, 1)), 'constant',constant_values=float('inf'))
      # find top scores
      score_ranked = np.sort(score_set, axis=None)
      score_ranked[score_ranked == float('inf')] = 0
      score_ranked = np.trim_zeros(score_ranked)
      
##      print("score ranked: " + __name__, score_ranked.shape)
      idx_ranked = np.argsort(score_set, axis=None)
      updated_beam_set = []
      for new_beam_rank in range(np.min((len(score_ranked), 10))): # args.beam_size=10
        total_idx = np.unravel_index(idx_ranked[new_beam_rank], score_set.shape)
        prev_beam_rank = total_idx[0]
        cluster_seq = total_idx[1:]
        updated_beam_state = _update_beam_state(beam_set[prev_beam_rank], look_ahead_seq, cluster_seq)
        updated_beam_set.append(updated_beam_state)
      beam_set = updated_beam_set
    predicted_cluster_id = beam_set[0].trace[-test_sequence_length:]
    
    
    return predicted_cluster_id

def predict(test_sequences, args):

    print("Predict method")
    # check type
    if isinstance(test_sequences, np.ndarray):
      return predict_single(test_sequences, args)
    if isinstance(test_sequences, list):
      return [predict_single(test_sequence, args) for test_sequence in test_sequences]
    raise TypeError('test_sequences should be either a list or numpy array.')


SAVED_MODEL_NAME = 'C:/Users/User/Speaker-Diarization/pretrained/saved_model.uisrnn_benchmark'
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
observation_dim = 512

model_args, _, inference_args = uisrnn.parse_arguments()
args = uisrnn.parse_arguments()
model_args.observation_dim = observation_dim
rnn_init_hidden = nn.Parameter(torch.zeros(model_args.rnn_depth, 1, model_args.rnn_hidden_size).to(device))

var_dict = load(SAVED_MODEL_NAME)
rnn_init_hidden = nn.Parameter(torch.from_numpy(var_dict['rnn_init_hidden']).to(device))
_INITIAL_SIGMA2_VALUE = 0.1
estimate_sigma2 = (model_args.sigma2 is None)
estimate_transition_bias = (model_args.transition_bias is None)
sigma2 = nn.Parameter(torch.from_numpy(var_dict['sigma2']).to(device))
transition_bias = float(var_dict['transition_bias'])
transition_bias_denominator = float(var_dict['transition_bias_denominator'])
crp_alpha = float(var_dict['crp_alpha'])
rnn_model = CoreRNN(model_args.observation_dim, model_args.rnn_hidden_size, model_args.rnn_depth, model_args.observation_dim, model_args.rnn_dropout).to(device)
rnn_model.load_state_dict(var_dict['rnn_state_dict'])




batch_size = 1
x = np.ones((batch_size, 512)).tolist()
x = torch.randn(1, batch_size, 512)
x = np.zeros((40,512))

model = torch.jit.script(predict(x,model_args))
#model.load(SAVED_MODEL_NAME)
torch.onnx.export(uisrnnModel, x, "uisRNN_core.onnx", do_constant_folding=False, export_params=True, input_names = ['input'], output_names = ['output'], dynamic_axes={'input' : {0 : 'batch_size', 1:'utterance_size'}})
print("onnx model exported")
</code></pre>
<p>I also tried other tweaks, such as the use of decorator <span class="mention">@torch.jit.script</span> at the very top of my script. This produces the following error:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "C:\Users\User\Speaker-Diarization\uisrnn_plain_functions_old.py", line 15, in &lt;module&gt;
    class CoreRNN(nn.Module):
  File "C:\Users\User\Python\lib\site-packages\torch\jit\_script.py", line 909, in script
    " pass an instance instead".format(obj)
RuntimeError: Type '&lt;class '__main__.CoreRNN'&gt;' cannot be compiled since it inherits from nn.Module, pass an instance instead
</code></pre>
<p>Finally, if I try to export without the use of jit, I get this error:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "C:\Users\User\Speaker-Diarization\uisrnn_plain_functions_old.py", line 271, in &lt;module&gt;
    torch.onnx.export(model, x, "uisRNN_core.onnx", do_constant_folding=False, export_params=True, input_names = ['input'], output_names = ['output'], dynamic_axes={'input' : {0 : 'batch_size', 1:'utterance_size'}})
  File "C:\Users\User\Python\lib\site-packages\torch\onnx\__init__.py", line 230, in export
    custom_opsets, enable_onnx_checker, use_external_data_format)
  File "C:\Users\User\Python\lib\site-packages\torch\onnx\utils.py", line 91, in export
    use_external_data_format=use_external_data_format)
  File "C:\Users\User\Python\lib\site-packages\torch\onnx\utils.py", line 618, in _export
    with select_model_mode_for_export(model, training):
  File "C:\Users\User\Python\lib\contextlib.py", line 112, in __enter__
    return next(self.gen)
  File "C:\Users\User\Python\lib\site-packages\torch\onnx\utils.py", line 35, in select_model_mode_for_export
    is_originally_training = model.training
AttributeError: 'list' object has no attribute 'training'
</code></pre>
<p>What should I do to export this model? Every plan fell through. I’ve tried all possible combinations of placing decorators and using the <code>torch.jit.script</code> and <code>torch.jit.trace</code> functions. I’ve also attempted to pass in an instance of the model calling model= CoreRNN(), but to no avail. Every contrived experiment from PyTorch’s website I have attempted to recreate works fine. ΙAre there any more advanced tutorials on this? Any help is greatly appreciated.</p>