Torch JIT trace problem

jit
1

The code is as follows:

class DemoTargetObjectFeatureProcessor(torch.nn.Module):
    def __init__(self):
        super(DemoTargetObjectFeatureProcessor, self).__init__()

    def forward(self, target_object_inputs):
        target_object_size = target_object_inputs[:, 0] * target_object_inputs[:, 1]
        target_object_size = target_object_size.unsqueeze(-1)
        return torch.cat((target_object_inputs, target_object_size), dim=1)

    @torch.jit.export
    def forward_1(self, target_object_inputs_dict):
        target_object_inputs = target_object_inputs_dict['a']
        return self.forward(target_object_inputs)

target_object_inputs = torch.tensor([[1,1], [2,2]])
target_object_inputs_dict = {
    "a": torch.tensor([[1,1], [2,2]])
}

fp = DemoTargetObjectFeatureProcessor()
fp.forward(target_object_inputs)
fp.forward_1(target_object_inputs_dict)

module = torch.jit.trace(fp, target_object_inputs)
module = torch.jit.trace(fp.forward, target_object_inputs)

module = torch.jit.trace(fp.forward_1, target_object_inputs_dict)

I’m trying to pass a named dictionary of tensors into the TorchScript module from “module = torch.jit.trace(fp.forward_1, target_object_inputs_dict)”, but I will get the error as follows:

Traceback (most recent call last):
  File "test_trace.py", line 50, in <module>
    module = torch.jit.trace(fp.forward_1, target_object_inputs_dict)
  File "/home/lei.chen/.pyenv/versions/3.6.6/lib/python3.6/site-packages/torch/jit/__init__.py", line 893, in trace
    raise AttributeError("trace doesn't support compiling individual module's functions.\n"
AttributeError: trace doesn't support compiling individual module's functions.
Please use trace_module

If i’m trying to use

module = torch.jit.script(fp)
module.forward_1(target_object_inputs_dict)

I will get the follow error:

RuntimeError:
Unsupported operation: indexing tensor with unsupported index type 'str'. Only ints, slices, and tensors are supported:
  File "test_trace.py", line 36
    @torch.jit.export
    def forward_1(self, target_object_inputs_dict):
        target_object_inputs = target_object_inputs_dict['a']
                               ~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
        return self.forward(target_object_inputs)

How can we pass a named dictionary of tensors into TorchScript? I want to run model inference in C++ backend using the exported TorchScript. From the PyTorch unit tests https://github.com/pytorch/pytorch/blob/5136ed0e44c65cb3747a1f22f77ccf09d54c125c/test/cpp/api/jit.cpp#L73, it seems that we can pass a c10::Dict into TorchScript module.

2

Any TorchScript expert can help here? I also find this issue in PyTorch github page: https://github.com/pytorch/pytorch/issues/16847. It’s claimed that “Closing as we support this in the tracer now.”, but many ppls experienced the error from dictionary as input tensors.

3

Hi, if you read the error, it infers the input as a tensor. You need to explicitly specify it’s a dict with annotation

from typing import Dict
import torch

def forward_1(self, target_object_inputs_dict: Dict[str, torch.Tensor]):