torch.jit.script(module) vs @torch.jit.script decorator

user_4974732 0
Viewed 2141

Why is adding the decorator "@torch.jit.script" results in an error, while I can call torch.jit.script on that module, e.g. this fails:

import torch
    
@torch.jit.script
class MyCell(torch.nn.Module):
    def __init__(self):
        super(MyCell, self).__init__()
        self.linear = torch.nn.Linear(4, 4)

    def forward(self, x, h):
        new_h = torch.tanh(self.linear(x) + h)
        return new_h, new_h
    
my_cell = MyCell()
x, h = torch.rand(3, 4), torch.rand(3, 4)
traced_cell = torch.jit.script(my_cell, (x, h))
print(traced_cell)
traced_cell(x, h)

"C:\Users\Administrator\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.8_qbz5n2kfra8p0\LocalCache\local-packages\Python38\site-packages\torch\jit\__init__.py", line 1262, in script
    raise RuntimeError("Type '{}' cannot be compiled since it inherits"
RuntimeError: Type '<class '__main__.MyCell'>' cannot be compiled since it inherits from nn.Module, pass an instance instead

While the following code works well:

class MyCell(torch.nn.Module):
    def __init__(self):
        super(MyCell, self).__init__()
        self.linear = torch.nn.Linear(4, 4)
    
    def forward(self, x, h):
        new_h = torch.tanh(self.linear(x) + h)
        return new_h, new_h
    
my_cell = MyCell()
x, h = torch.rand(3, 4), torch.rand(3, 4)
traced_cell = torch.jit.script(my_cell, (x, h))
print(traced_cell)
traced_cell(x, h)

This question is also featured on PyTorch forums.

torchscript pytorch
1 Answers

Reason for your error is here, this bulletpoint precisely:

No support for inheritance or any other polymorphism strategy, except for inheriting from object to specify a new-style class.

Also, as stated at the top:

TorchScript class support is experimental. Currently it is best suited for simple record-like types (think a NamedTuple with methods attached).

Currently, it's purpose is for simple Python classes (see other points in the link I've provided) and functions, see link I've provided for more information.

You can also check torch.jit.script source code to get a better grasp of how it works.

From what it seems, when you pass an instance, all attributes which should be preserved are recursively parsed (source). You can follow this function along (quite commented, but too long for an answer, see here), though exact reason why this is the case (and why it was designed this way) is beyond my knowledge (so hopefully someone with expertise in torch.jit's inner workings will speak more about it).

user_108864200
Related