Using TorschScript to save a model with multiple heads

giii: <p>My goal is to serialize a pytorch trained model an load it in an environment where the original class defining the neural network is not available. To achieve that, I decided to use TorchScript since it seems the only possible way.</p>
<p>I have a multi-task model (type nn.Module) built using a body common to every task (also nn.Module, just a few linear layers) and a set of linear head models, one per task.<br>
I store the head models in a dictionary <code>Dict[int, nn.Module]</code> called <code>_task_head_models</code> and I created an ad-hoc forward method in my module class to select the right head at prediction time:</p>
<pre><code class="lang-auto">    def forward(self, x: torch.Tensor, task_id: int = 0) -&gt; torch.Tensor:
        if task_id not in self._task_head_models.keys():
            raise ValueError(
                f"The task id {task_id} is not valid. Valid task ids are {self._task_head_models.keys()}."
            )

        return self._task_head_models[task_id](self._model(x))
</code></pre>
<p>This works fine until I am not trying to serialize it using torchscript.<br>
When I try <code>torch.jit.script(mymodule)</code>, I get:</p>
<pre><code class="lang-auto">Module 'MyModule' has no attribute '_task_head_models' (This attribute exists on the Python module, but we failed to convert Python type: 'dict' to a TorchScript type. Cannot infer concrete type of torch.nn.Module. Its type was inferred; try adding a type annotation for the attribute.)
</code></pre>
<p>Something that seems off, is that my module contains a <code>Dict</code>, not a <code>dict</code> as mentioned in the error message. Forgetting that for a second, it’s still unclear why this is happening. Dictionaries seems to be supported in the language reference: <a href="https://docs.w3cub.com/pytorch/jit_language_reference.html" class="inline-onebox" rel="noopener nofollow ugc">TorchScript Language Reference - PyTorch - W3cubDocs</a></p>
<p>I also tried to use ModuleDict instead of Dict (changing the key type to str) but that doesn’t seem to work either: <code>Unable to extract string literal index. ModuleDict indexing is only supported with string literals. Enumeration of ModuleDict is supported, e.g. 'for k, v in self.items(): ...':</code></p>

Cristian_Regep: <p>I’ve encountered the same problem and after reading this post <a href="https://github.com/pytorch/pytorch/issues/47496" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/issues/47496</a> and seeing this code in their tests <a href="https://github.com/pytorch/pytorch/blob/2e57b1e8479bb09ea1b0a2de059c1322e6782fdb/test/jit/test_module_containers.py#L431" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/2e57b1e8479bb09ea1b0a2de059c1322e6782fdb/test/jit/test_module_containers.py#L431</a> I figured it out. To get non-static indexing to work you need to define an interface class with the jit decorator. I’m writing it here so others don’t have to go through this pain</p>
<p>Let’s say the heads are of type Head</p>
<pre><code class="lang-python">class Head(nn.Module):
....
  def forward(self, my_x: Tensor)-&gt;Tensor:
    return self.do_something(x)
</code></pre>
<p>You need to write an interface with the exact same signature (identical names!) and provide the jit.interface decorator</p>
<pre><code class="lang-python">@torch.jit.interface
class OutputHeadInterface(nn.Module):
    def forward(self, my_x: Tensor) -&gt; Tensor: 
      pass
</code></pre>
<p>Then, in the forward of your model, using your code:</p>
<pre><code class="lang-python">class MyModule
.....
def forward(self, x: torch.Tensor, task_id: str = "0") -&gt; torch.Tensor:
        if task_id not in self._task_head_models.keys():
            raise ValueError(
                f"The task id {task_id} is not valid. Valid task ids are {self._task_head_models.keys()}."
            )

        output_head: OutputHeadInterface = self._task_head_models[task_id]
        return output_head.forward(self._model(x))
</code></pre>
<p>I’ve changed the task_id to a string</p>