pytorch - 使用 _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

作为 [ https://stackoverflow.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