How to make jit script works with a custom RNN module (renamed)

jit
1

Hi,
I created an issue at:

However, I am thinking maybe here is the right place to ask for help. Thus I copied my ask here:

This is regarding sample at:
https://pytorch.org/blog/the-road-to-1_0/:

from torch.jit import script

@script
def rnn_loop(x):
hidden = None
for x_t in x.split(1):
x, hidden = model(x, hidden)
return x

I cannot make it to work. Here is my code:

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

def test_ScriptModelRNN():
class SimpleRNNCell(nn.Module):
def  **init** (self, input_size, hidden_size):
super(SimpleRNNCell, self). **init** ()
self.linear_h = nn.Linear(input_size, hidden_size)

    def forward(self, inp, h_0):
        h = self.linear_h(inp)
        return h + h_0, h

with torch.no_grad():
    sequence_len, input_size, hidden_size = 4, 3, 2
    model = SimpleRNNCell(input_size, hidden_size)
    hidden = torch.zeros(1, hidden_size)

    # # test cell
    # cell_input = torch.randn(input_size)
    # cell_output, hidden = model(cell_input, hidden)
    # import pdb; pdb.set_trace()
    # #

    @torch.jit.script
    def rnn_loop(x):
        hidden = None
        for x_t in x.split(1):
            x, hidden = model(x_t, hidden)
        return x
        
    input = torch.randn(sequence_len, input_size)
    output = rnn_loop(input)

I am getting:

Exception has occurred: RuntimeError

for operator (Tensor 0, Tensor 1) -> (Tensor, Tensor):
expected a value of type Tensor for argument '1' but found Tensor?
@torch.jit.script
def rnn_loop(x):
hidden = None
for x_t in x.split(1):
x, hidden = model(x_t, hidden)
~~~~~~ <--- HERE
return x
:
@torch.jit.script
def rnn_loop(x):
hidden = None
for x_t in x.split(1):
x, hidden = model(x_t, hidden)
~~~~~ <--- HERE
return x
File "/home/liqun/pytorch/torch/jit/ **init** .py", line 751, in script
_jit_script_compile(mod, ast, _rcb, get_default_args(obj))
File "/home/liqun/Untitled Folder/test_onnx_export.py", line 218, in test_ScriptModelRNN
@torch.jit.script
File "/home/liqun/Untitled Folder/test_onnx_export.py", line 282, in 
test_ScriptModelRNN()
File "/home/liqun/.conda/envs/py36/lib/python3.6/runpy.py", line 85, in _run_code
exec(code, run_globals)
File "/home/liqun/.conda/envs/py36/lib/python3.6/runpy.py", line 96, in _run_module_code
mod_name, mod_spec, pkg_name, script_name)
File "/home/liqun/.conda/envs/py36/lib/python3.6/runpy.py", line 263, in run_path
pkg_name=pkg_name, script_name=fname)
2

Thanks for the question, @liqun.

The problem looks like you’re passing None as the hidden:

def rnn_loop(x):
hidden = None
**for** x_t **in** x.split(1):
x, hidden = model(x_t, hidden)

However, you’re not handling None hidden values in your model.

Please see the following file https://github.com/pytorch/benchmark/blob/master/rnns/fastrnns/custom_lstms.py for some samples of custom RNNs written with TorchScript.

3

Thanks Richard!
The link is very helpful. I will take a close look.
BTW, if I assign a tensor value instead of None to the hidden, like this:
hidden = torch.zeros(1, hidden_size)
I am getting:
python value of type ‘int’ cannot be used as a value:
@torch.jit.script
def rnn_loop(x):
hidden = torch.zeros(1, hidden_size)
~~~~~~~~~~~ <— HERE
for x_t in x.split(1):
x, hidden = model(x_t, hidden)
return x
File “/home/liqun/pytorch/torch/jit/init.py”, line 751, in script
_jit_script_compile(mod, ast, _rcb, get_default_args(obj))

Actually this was the error I have been getting. The None assignment was just to workaround this error to see how far I can go.

And if I assign a tensor value to the hidden outside of the loop, I am getting:
python value of type ‘Tensor’ cannot be used as a value:
@torch.jit.script
def rnn_loop(x):
#hidden = torch.zeros(1, hidden_size)
for x_t in x.split(1):
x, hidden = model(x_t, hidden)
~~~~~~ <— HERE

Thanks again,
Liqun