torch.nn.Sequential of designed blocks problem in giving inputs

Question

I have designed a class that is a block of a network and its forward has three inputs: x, logdet, reverse, and has two outputs. for example, everything is normal when I call this class and use it, like:

x = torch.Tensor(np.random.rand(2, 48, 8, 8))
net = Block(inp = 48, oup = 48, mid_channels=48, ksize=3, stride=1, group = 3)
a, _ = net(x, reverse = False)

But when I want to use it by Sequential(because I need multi blocks after each other), the problem occurs like this:

x = torch.Tensor(np.random.rand(2, 48, 8, 8))
conv1_network = nn.Sequential(
    Block(inp = 48, oup = 48, mid_channels=48, ksize=3, stride=1, group = 3)
        )
conv1_network(x, reverse = False)

My error is: TypeError: forward() got an unexpected keyword argument 'reverse' And it is not normal because I have reverse in my inputs of forward in Block as we see in the first part. I'm looking forward to finding a way to attach some Blocks to each other for example this is a block

class Block(nn.Module):
    def __init__(self, num_channels):
        super(InvConv, self).__init__()
        self.num_channels = num_channels

        # Initialize with a random orthogonal matrix
        w_init = np.random.randn(num_channels, num_channels)
        w_init = np.linalg.qr(w_init)[0].astype(np.float32)
        self.weight = nn.Parameter(torch.from_numpy(w_init))

    def forward(self, x, logdet, reverse=False):
        ldj = torch.slogdet(self.weight)[1] * x.size(2) * x.size(3)

        if reverse:
            weight = torch.inverse(self.weight.double()).float()
            logdet = logdet - ldj
        else:
            weight = self.weight
            logdet = logdet + ldj

        weight = weight.view(self.num_channels, self.num_channels, 1, 1)
        z = F.conv2d(x, weight)

        return z, logdet

And my purpose is to attach multi Blocks to each other in Sequential in a for(because I can't use the same Block in my work, I need different convolutions for making a deep network)

features = []
for i in range(10):
   self.features.append(Block(num_channels = 48))

and then I want to use them like this

self.features(x, logdet = 0, reverse = False)

Answer 1

You indicated that your Block nn.Module had a reverse option. However nn.Sequential doesn't, so conv1_network(x, reverse=False) is not valid because conv1_network is not a Block.

By default, you can't pass kwargs to layers inside a nn.Sequential. You can however inherit from nn.Sequential and do it yourself. Something like:

class BlockSequence(nn.Sequential):
    def forward(self, input, **kwargs):
        for module in self:
            options = kwargs if isinstance(module, Block) else {}
            input = module(input, **options)
        return input

This way, you can create a sequence containing Blocks (and optionally non-Block modules as well):

>>> blocks = []
>>> for i in range(10):
...     self.blocks.append(Block(num_channels=48))

>>> blocks = BlockSequence(*blocks)

Then you will be able to call blocks with the reverse keyword argument, which be relayed to every potential Block child module when called:

>>> blocks(x, logdet=0, reverse=False)