使用TorchScript类作为pytorch模块中的成员

Question

我正在尝试让一些现有的pytorch模型支持TorchScript jit编译器，但是我遇到了非原语类型的成员问题。

这个小例子说明了这个问题：

import torch

@torch.jit.script
class Factory(object):
    def __init__(self):
        pass

    def create(self, x: float) -> torch.Tensor:
        return torch.tensor([x])

class Foo(torch.nn.Module):
    def __init__(self):
        super(Foo, self).__init__()
        self.factory: Factory = Factory()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.factory.create(0)

mod = torch.jit.script(Foo())

运行时，jit编译器会给出错误

RuntimeError:
module has no attribute 'factory':
at example.py:17:15
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.factory.create(0)
               ~~~~~~~~~~~~ <--- HERE

我已经测试了forward方法中的jit可以使用Factory类，但是当我将它存储为成员时，它不会承认它。为什么会这样呢？有没有办法让jit编译器将这种成员保存到编译后的模块中？

Answer 1

这是PyTorch中的一个bug，在你发布你的问题：https://discuss.pytorch.org/t/jit-scripted-attributes-inside-module/60645，https://github.com/pytorch/pytorch/issues/27495后不久就解决了。

更新PyTorch应该可以解决这个问题。