Torchscripting a module with _ConvNd in forward - pytorch

Question

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.

Answer 1

As alternative to [https://stackoverflow.com/users/6210807/kharshit] suggestion, you can define network functional way:
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

Answer 2

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