Torchscripting a module with _ConvNd in forward

I am using PyTorch 1.4 and need to export a model with convolutions inside a loop in 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())

This gives the following error:

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

I have also tried pre-defining the conv's then putting them in a list inside __init__, but such a type is not allowed by 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())

This instead gives:

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

So how to export this module? Background: I am exporting Mixed-scale Dense Networks (source) to TorchScript; while nn.Sequential may work for this simplified case, practically I need to convolve with all the historical convolution outputs in each iteration, which is more than chaining the layers.

Solution

You can use nn.ModuleList() in the following way.

Also, note that you can't subscript nn.ModuleList currently probably due to a bug as mentioned in issue#16123, but use the workaround as mentioned below.

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)
)