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:
Traceback (most recent call last):
  File "C:\Users\User\Speaker-Diarization\uisrnn_plain_functions_old.py", line 269, in <module>
    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

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:

Traceback (most recent call last):
  File "C:\Users\User\Speaker-Diarization\uisrnn_plain_functions_old.py", line 269, in <module>
    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

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:
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 >= 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 > len(new_beam_state.mean_set):  # invalid trace
        new_beam_state.neg_likelihood = float('inf')
        break
      elif cluster < 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")

I also tried other tweaks, such as the use of decorator @torch.jit.script at the very top of my script. This produces the following error:
Traceback (most recent call last):
  File "C:\Users\User\Speaker-Diarization\uisrnn_plain_functions_old.py", line 15, in <module>
    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 '<class '__main__.CoreRNN'>' cannot be compiled since it inherits from nn.Module, pass an instance instead

Finally, if I try to export without the use of jit, I get this error:
Traceback (most recent call last):
  File "C:\Users\User\Speaker-Diarization\uisrnn_plain_functions_old.py", line 271, in <module>
    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'

What should I do to export this model? Every plan fell through. I’ve tried all possible combinations of placing decorators and using the torch.jit.script and torch.jit.trace 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.