Bayesian Model output size changes with input image size

[dipayandas97]

Congratulations on this wonderful achievement of the paper. I am trying to implement the Bayesian CNN for a project of mine, a simple binary classification. Here is the code to define the model:

import torch.nn as nn
import math
from layers import BBB_Linear, BBB_Conv2d
from layers import BBB_LRT_Linear, BBB_LRT_Conv2d
from layers import FlattenLayer, ModuleWrapper

class get_model(ModuleWrapper):
    def __init__(self, outputs, inputs, priors=None, layer_type='bbb', activation_type='softplus'):
        super(get_model, self).__init__()

        self.num_classes = outputs
        self.layer_type = layer_type
        self.priors = priors

        if layer_type=='lrt':
            BBBLinear = BBB_LRT_Linear
            BBBConv2d = BBB_LRT_Conv2d
        elif layer_type=='bbb':
            BBBLinear = BBB_Linear
            BBBConv2d = BBB_Conv2d
        else:
            raise ValueError("Undefined layer_type")

        if activation_type=='softplus':
            self.act = nn.Softplus
        elif activation_type=='relu':
            self.act = nn.ReLU
        else:
            raise ValueError("Only softplus or relu supported")

        self.conv1 = BBBConv2d(inputs, 64, 11, stride=4, padding=5, bias=True, priors=self.priors)
        self.act1 = self.act()
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.conv2 = BBBConv2d(64, 192, 5, padding=2, bias=True, priors=self.priors)
        self.act2 = self.act()
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.conv3 = BBBConv2d(192, 384, 3, padding=1, bias=True, priors=self.priors)
        self.act3 = self.act()

        self.conv4 = BBBConv2d(384, 256, 3, padding=1, bias=True, priors=self.priors)
        self.act4 = self.act()

        self.conv5 = BBBConv2d(256, 128, 3, padding=1, bias=True, priors=self.priors)
        self.act5 = self.act()
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.flatten = FlattenLayer(1 * 1 * 128)

        self.fc1 = BBBLinear(128, 64, bias=True, priors=self.priors)
        self.act6 = self.act()

        self.fc2 = BBBLinear(64, 32, bias=True, priors=self.priors)
        self.act7 = self.act()

        self.fc3 = BBBLinear(32, outputs, bias=True, priors=self.priors)

model = get_model(outputs=2, inputs=1).to(device)

However, when i input a random input like,

model(torch.randn(1,1,64,64).to(device))[0].size()

the ouput has a shape (4,2) whereas it should be (1,2). The peculiarity seems like the model outputs (1,2) only when it gets an input of shape (1,32,32), like the MNIST dataset. However I could not find where this fixation of 32x32 image size is done.

Looking forward to your support! Thanks

[dipayandas97]

The issue was solved. The arguments of Flatten(11128) needs to be adjusted according to input image size.

[Eastonboy99]

Hello @dipayandas97,

I am also trying to use this for a binary classification. As I am very new to deep learning I am a bit confused and was hoping you should share what you changed the arguments of flatten to?

Thanks in advance