CODEWITHSUNDEEP
pythonpytorchjittorchscript
Chuanhua Yang
Chuanhua Yang

Reputation: 133

how to use custom python object in torchscript

I am ready to convert a pytorch module to ScriptModule and then load it in c++，but I am blocked by this error This attribute exists on the Python module, but we failed to convert Python type: 'Vocab' to a TorchScript type, the Vocab is a python object I define. the demo code is here:

import torch
class Vocab(object):
    def __init__(self, name):
        self.name = name

    def show(self):
        print("dict:" + self.name)

class Model(torch.nn.Module):
    def __init__(self, ):
        super(Model, self).__init__()
        self.layers = torch.nn.Linear(2, 3)
        self.encoder = 4
        self.vocab = Vocab("vocab")

    def forward(self, x):
        name = self.vocab.name
        print("forward show encoder:" + str(self.encoder))
        print("vocab:" + name)
        enc_hidden = []
        step = len(x) // 2
        for i in range(step):
            enc_hidden.append((x[2*i] + x[2*i + 1])/2)
        enc_hidden = torch.stack(enc_hidden, 0)
        enc_hidden = self.__show(enc_hidden)
        return self.layers(enc_hidden)

    @torch.jit.export
    def __show(self, x):
        return x + 1

model = Model()
data = torch.randn(10, 2)
script_model = torch.jit.script(model)
print(script_model)
r1 = model(data)
print(r1)

the error msg:

Traceback (most recent call last):
  File "/mnt/d/python_projects/pytorch_deploy/model4.py", line 47, in <module>
    script_model = torch.jit.script(model)
  File "/mnt/d/anaconda3/lib/python3.6/site-packages/torch/jit/__init__.py", line 1261, in script
    return torch.jit._recursive.create_script_module(obj, torch.jit._recursive.infer_methods_to_compile)
  File "/mnt/d/anaconda3/lib/python3.6/site-packages/torch/jit/_recursive.py", line 305, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/mnt/d/anaconda3/lib/python3.6/site-packages/torch/jit/_recursive.py", line 361, in create_script_module_impl
    create_methods_from_stubs(concrete_type, stubs)
  File "/mnt/d/anaconda3/lib/python3.6/site-packages/torch/jit/_recursive.py", line 279, in create_methods_from_stubs
    concrete_type._create_methods(defs, rcbs, defaults)
RuntimeError: 
  Module 'Model' has no attribute 'vocab' (This attribute exists on the Python module, but we failed to convert Python type: 'Vocab' to a TorchScript type.):
  File "/mnt/d/python_projects/pytorch_deploy/model4.py", line 26
  def forward(self, x):
    name = self.vocab.name
           ~~~~~~~~~~ <--- HERE
    print("forward show encoder:" + str(self.encoder))
    print("vocab:" + name)

so how can I use my own python object in torchscript?

Upvotes: 5

Views: 5897

Answers (1)

Szymon Maszke
Szymon Maszke

Reputation: 24691

You have to annotate your Vocab with torchscript.jit like this:

@torch.jit.script
class Vocab(object):
    def __init__(self, name: str):
        self.name = name

    def show(self):
        print("dict:" + self.name)

Also note specification name: str as it's also needed for torchscript to infer it's type (PyTorch supports >=Python3.6 type annotations, you could use a comment as well, but it's way less clear).

Please see Torchscript classes and Default Types and other related torchscript info over there.

Upvotes: 10

Related Questions