使用 _ConvNd 对模块进行 Torchscripting

Question

我正在使用 PyTorch 1.4，需要在 forward 的循环中导出带有卷积的模型：

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

    def forward(self, x):
        for i in range(5):
            conv = torch.nn.Conv1d(1, 1, 2*i+3)
            x = torch.nn.Relu()(conv(x))
        return x


torch.jit.script(MyCell())

这会产生以下错误：

RuntimeError: 
Arguments for call are not valid.
The following variants are available:

  _single(float[1] x) -> (float[]):
  Expected a value of type 'List[float]' for argument 'x' but instead found type 'Tensor'.

  _single(int[1] x) -> (int[]):
  Expected a value of type 'List[int]' for argument 'x' but instead found type 'Tensor'.

The original call is:
  File "***/torch/nn/modules/conv.py", line 187
                 padding=0, dilation=1, groups=1,
                 bias=True, padding_mode='zeros'):
        kernel_size = _single(kernel_size)
                      ~~~~~~~ <--- HERE
        stride = _single(stride)
        padding = _single(padding)
'Conv1d.__init__' is being compiled since it was called from 'Conv1d'
  File "***", line ***
    def forward(self, x):
        for _ in range(5):
            conv = torch.nn.Conv1d(1, 1, 2*i+3)
                   ~~~~~~~~~~~~~~~ <--- HERE
            x = torch.nn.Relu()(conv(x))
        return x
'Conv1d' is being compiled since it was called from 'MyCell.forward'
  File "***", line ***
    def forward(self, x, h):
        for _ in range(5):
            conv = torch.nn.Conv1d(1, 1, 2*i+3)
            ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
            x = torch.nn.Relu()(conv(x))
        return x

我也尝试过预定义 conv 然后将它们放在 __init__ 内的列表中，但 TorchScript 不允许这样的类型：

class MyCell(torch.nn.Module):
    def __init__(self):
        super(MyCell, self).__init__()
        self.conv = [torch.nn.Conv1d(1, 1, 2*i+3) for i in range(5)]

    def forward(self, x):
        for i in range(len(self.conv)):
            x = torch.nn.Relu()(self.conv[i](x))
        return x


torch.jit.script(MyCell())

这反而给出：

RuntimeError: 
Module 'MyCell' has no attribute 'conv' (This attribute exists on the Python module, but we failed to convert Python type: 'list' to a TorchScript type.):
  File "***", line ***
    def forward(self, x):
        for i in range(len(self.conv)):
                           ~~~~~~~~~ <--- HERE
            x = torch.nn.Relu()(self.conv[i](x))
        return x

那么如何导出这个模块呢？背景：我正在将 Mixed-scale Dense Networks (source) 导出到 TorchScript；虽然 nn.Sequential 可能适用于这种简化的情况，但实际上我需要在每次迭代中与所有历史卷积输出进行卷积，这不仅仅是链接层。

Answer 1

您可以通过以下方式使用nn.ModuleList()。

另外请注意，您目前无法下标 nn.ModuleList，这可能是由于 issue#16123 中提到的错误，但请使用下面提到的解决方法。

class MyCell(nn.Module):
    def __init__(self):
        super(MyCell, self).__init__()
        self.conv = nn.ModuleList([torch.nn.Conv1d(1, 1, 2*i+3) for i in range(5)])
        self.relu = nn.ReLU()

    def forward(self, x):
        for mod in self.conv:
            x = self.relu(mod(x))
        return x

>>> torch.jit.script(MyCell())
RecursiveScriptModule(
  original_name=MyCell
  (conv): RecursiveScriptModule(
    original_name=ModuleList
    (0): RecursiveScriptModule(original_name=Conv1d)
    (1): RecursiveScriptModule(original_name=Conv1d)
    (2): RecursiveScriptModule(original_name=Conv1d)
    (3): RecursiveScriptModule(original_name=Conv1d)
    (4): RecursiveScriptModule(original_name=Conv1d)
  )
  (relu): RecursiveScriptModule(original_name=ReLU)
)

Answer 2

作为 [https://whosebug.com/users/6210807/kharshit] 建议的替代方案，您可以定义网络功能方式：

class MyCell(torch.nn.Module):
    def __init__(self):
        super(MyCell, self).__init__()
        self.w = []
        for i in range(5):
            self.w.append( torch.Tensor( 1, 1, 2*i+3 ) )
            # init w[i] here, maybe make it "requires grad" 

    def forward(self, x):
        for i in range(5):
            x = torch.nn.functional.conv1d( x, self.w[i] )
            x = torch.nn.functional.relu( x )
        return x

使用 _ConvNd 对模块进行 Torchscripting

Torchscripting a module with _ConvNd in forward

pytorch

torchscript