Get the CNN layer output size in def init PyTorch

Viewed 996

When defining our model architecture in PyTorch, we need to specify the size of CNN output layer to feed into the nn.Linear layer. How can we find the size of this layer in the def __init__ function (not in def forward())

class model(nn.Module):

    def __init__(self,word_count,img_channel,n_out):
        super(multimodal,self).__init__()

        # CNN image encoding hyperparameters
        conv1_channel_out = 8
        conv1_kernel = 5

        pool1_size = 2

        conv2_channel_out = 16
        conv2_kernel = 16

        pool2_size = 2

        conv3_channel_out = 32
        conv3_kernel = 4

        dropout_rate = 0.1

        cnn_fc_out = 512

        comb_fc1_out = 512
        comb_fc2_out = 128

        # FNN text encoding hyperparameters
        text_fc1_out = 4096
        text_fc2_out = 512


        # Text encoding
        self.text_fc1 = nn.Linear(word_count, text_fc1_out)
        self.text_fc2 = nn.Linear(text_fc1_out, text_fc2_out)


        # Image encoding
        self.conv1 = nn.Conv2d(img_channel, conv1_channel_out, conv1_kernel)
        self.max_pool1 = nn.MaxPool2d(pool1_size)
        self.conv2 = nn.Conv2d(conv1_channel_out, conv2_channel_out, conv2_kernel)
        self.max_pool2 = nn.MaxPool2d(pool2_size)
        self.conv3 = nn.Conv2d(conv2_channel_out, conv3_channel_out, conv3_kernel)

        self.cnn_dropout = nn.Dropout(dropout_rate)
        self.cnn_fc = nn.Linear(32*24*12, cnn_fc_out)

        #Concat layer
        concat_feat = cnn_fc_out + text_fc2_out
        self.combined_fc1 = nn.Linear(concat_feat, comb_fc1_out)
        self.combined_fc2 = nn.Linear(comb_fc1_out, comb_fc2_out)
        self.output_fc = nn.Linear(comb_fc2_out, n_out)


    def forward(self, text, img):
        # Image Encoding
        x = F.relu(self.conv1(img))
        x = self.max_pool1(x)
        x = F.relu(self.conv2(x))
        x = self.max_pool2(x)
        x = F.relu(self.conv3(x))
        x = x.view(-1, 32*24*12)
        x = self.cnn_dropout(x)
        img = F.relu(self.cnn_fc(x))

        # Text Encoding
        text = F.relu(self.text_fc1(text))
        text = F.relu(self.text_fc2(text))

        # Concat the features
        concat_inp = torch.cat((text, img), 1)
        out = F.relu(self.combined_fc1(concat_inp))
        out = F.relu(self.combined_fc2(out))
        return torch.sigmoid(self.output_fc(out))

If you see above, I define the size of CNN output layer as 322412 manually self.cnn_fc = nn.Linear(32*24*12, cnn_fc_out)

How can I avoid this? I know we might be able to call [model_name].[layer_name].in_features in def forward(), but not in def __init__()

2 Answers

I dont think there is a specific way to do that. You would have to run a sample (you can just use x = torch.rand((1, C, W, H)) for testing) and then in forward print out the shape of the conv layer right before your linear layer, then you memorize that number and hardcode it into init. Or you could use formulas to calculate the shape of a conv layer based on the dimensions of the input, kernel-size, padding, etc. Here is a thread about those formulas.

There is no general way to do this, since the input and output sizes are not fixed in a CNN. What you can output is the number of channels, but the module will accept and transform any image height X width dimensions (so long as hey are sufficiently large to produce results large enough for the next layer after unpadded convolutions and pooling etc).

Hence you cannot include this in init (naive to input, instantiation of object), only in forward (calculated upon seeing input).

Related