[Python] Using zip() on two nn.ModuleList

Is using two different nn.ModuleList() zipped lists correct to build the computational graph for training a neural net in PyTorch? nn.ModuleList is a wrapper around Python's list with a registration of a module for training.

I'm building a network which consists of 2x interchanging types of blocks in __init__:

def __init__(self, in_channels):
super().__init__()

self.encoder_conv_blocks = nn.ModuleList()
self.downsample_blocks = nn.ModuleList()

for out_channels in _FILTERS:
conv_block = _ConvBlock(in_channels, _CONV_BLOCK_LEN, _CONV_BLOCK_GROWTH_RATE)
downsample_block = _DownsampleBlock(conv_block.out_channels, out_channels)

self.encoder_conv_blocks.append(conv_block)
self.downsample_blocks.append(downsample_block)

in_channels = out_channels


later in forward, I'm zipping the layers, as I need the outputs of the first type of block later in skip connections:

def forward(self, x):
skip_connections = []

for conv_block, downsample_block in zip(self.encoder_conv_blocks,
self.downsample_blocks):
x = conv_block(x)
skip_connections.append(x)
x = downsample_block(x)


However when pritting the summary torchinfo, we can see that summary of the registered methods using 2x zipped nn.ModuleList looks different compared to the summary where one single nn.ModuleList was used. I suspect that this can cause issues for training and inference in the future.

zip(nn.ModuleList(), nn.ModuleList()):

========================================================================================================================
Layer (type:depth-idx) Input Shape Output Shape Param #
========================================================================================================================
MyNet [16, 4, 128, 256] [16, 3, 128, 256] --
├─ModuleList: 1-13 -- -- (recursive)
│ └─_ConvBlock: 2-1 [16, 4, 128, 256] [16, 84, 128, 256] 26,360
├─ModuleList: 1-14 -- -- (recursive)
│ └─_DownsampleBlock: 2-2 [16, 84, 128, 256] [16, 64, 64, 128] 48,448
├─ModuleList: 1-13 -- -- (recursive)
│ └─_ConvBlock: 2-3 [16, 64, 64, 128] [16, 144, 64, 128] 70,160
├─ModuleList: 1-14 -- -- (recursive)
│ └─_DownsampleBlock: 2-4 [16, 144, 64, 128] [16, 128, 32, 64] 166,016
├─ModuleList: 1-13 -- -- (recursive)
│ └─_ConvBlock: 2-5 [16, 128, 32, 64] [16, 208, 32, 64] 116,880
├─ModuleList: 1-14 -- -- (recursive)
│ └─_DownsampleBlock: 2-6 [16, 208, 32, 64] [16, 128, 16, 32] 239,744
├─ModuleList: 1-13 -- -- (recursive)
│ └─_ConvBlock: 2-7 [16, 128, 16, 32] [16, 208, 16, 32] 116,880
├─ModuleList: 1-14 -- -- (recursive)
│ └─_DownsampleBlock: 2-8 [16, 208, 16, 32] [16, 128, 8, 16] 239,744
├─ModuleList: 1-13 -- -- (recursive)
│ └─_ConvBlock: 2-9 [16, 128, 8, 16] [16, 208, 8, 16] 116,880
├─ModuleList: 1-14 -- -- (recursive)
│ └─_DownsampleBlock: 2-10 [16, 208, 8, 16] [16, 256, 4, 8] 479,488
├─ModuleList: 1-13 -- -- (recursive)
│ └─_ConvBlock: 2-11 [16, 256, 4, 8] [16, 336, 4, 8] 210,320
├─ModuleList: 1-14 -- -- (recursive)
│ └─_DownsampleBlock: 2-12 [16, 336, 4, 8] [16, 256, 2, 4] 774,400
├─ModuleList: 1-13 -- -- (recursive)
│ └─_ConvBlock: 2-13 [16, 256, 2, 4] [16, 336, 2, 4] 210,320
├─ModuleList: 1-14 -- -- (recursive)
│ └─_DownsampleBlock: 2-14 [16, 336, 2, 4] [16, 512, 1, 2] 1,548,800


single nn.ModuleList():

MyNet [16, 4, 128, 256] [16, 3, 128, 256] --
├─ModuleList: 1-1 -- -- --
│ └─_ConvBlock: 2-1 [16, 4, 128, 256] [16, 84, 128, 256] 26,360
│ └─_DownsampleBlock: 2-2 [16, 84, 128, 256] [16, 64, 64, 128] 48,448
│ └─_ConvBlock: 2-3 [16, 64, 64, 128] [16, 144, 64, 128] 70,160
│ └─_DownsampleBlock: 2-4 [16, 144, 64, 128] [16, 128, 32, 64] 166,016
│ └─_ConvBlock: 2-5 [16, 128, 32, 64] [16, 208, 32, 64] 116,880
│ └─_DownsampleBlock: 2-6 [16, 208, 32, 64] [16, 128, 16, 32] 239,744
│ └─_ConvBlock: 2-7 [16, 128, 16, 32] [16, 208, 16, 32] 116,880
│ └─_DownsampleBlock: 2-8 [16, 208, 16, 32] [16, 128, 8, 16] 239,744
│ └─_ConvBlock: 2-9 [16, 128, 8, 16] [16, 208, 8, 16] 116,880
│ └─_DownsampleBlock: 2-10 [16, 208, 8, 16] [16, 256, 4, 8] 479,488
│ └─_ConvBlock: 2-11 [16, 256, 4, 8] [16, 336, 4, 8] 210,320
│ └─_DownsampleBlock: 2-12 [16, 336, 4, 8] [16, 256, 2, 4] 774,400
│ └─_ConvBlock: 2-13 [16, 256, 2, 4] [16, 336, 2, 4] 210,320
│ └─_DownsampleBlock: 2-14 [16, 336, 2, 4] [16, 512, 1, 2] 1,548,800

Continue reading...