Problem adding AdaptiveAvgPool2d to finetune a pretrained densenet201

PapRazzi: <p>I want to add the following classifier to finetune a pretrained densenet201. The problem is it fails with a dimension error and I can’t figure out what’s up. Please help out</p>
<pre><code class="lang-python">class ClassifierNew(nn.Module):
    def __init__(self, inp = 1920, h1=1024, out = 7, d=0.35):
        super().__init__()
        self.ap = nn.AdaptiveAvgPool2d((1,1))
        self.mp = nn.AdaptiveMaxPool2d((1,1))
        self.fla = Flatten()
        self.bn0 = nn.BatchNorm1d(inp*2,eps=1e-05, momentum=0.1, affine=True)
        self.dropout0 = nn.Dropout(d)
        self.fc1 = nn.Linear(inp*2, h1)
        self.bn1 = nn.BatchNorm1d(h1,eps=1e-05, momentum=0.1, affine=True)
        self.dropout1 = nn.Dropout(d)
        self.fc2 = nn.Linear(h1, out)
        
    def forward(self, x):
        print(x.shape)
        ap = self.ap(x)
        mp = self.mp(x)
        x = torch.cat((ap,mp),dim=1)
        x = self.fla(x)
        x = self.bn0(x)
        x = self.dropout0(x)
        x = F.relu(self.fc1(x))
        x = self.bn1(x)
        x = self.dropout1(x)         
        x = self.fc2(x)
        
        return x
</code></pre>
<p>I get the following traceback</p>
<pre><code class="lang-auto">torch.Size([8, 1920])
torch.Size([8, 1920])
torch.Size([8, 1920])
torch.Size([8, 1920])
Traceback (most recent call last):
  File "main.py", line 130, in &lt;module&gt;
    model = trainroutines.train_validate_model(model, optimizer, criterion, config.n_epochs,data_loader_train , data_loader_val, config.minibatch_print_freq)
  File "/dccstor/sgdermatology/scripts/sg_fair_Newton6/trainroutines.py", line 55, in train_validate_model
    output = model(data)
  File "/u/nkinyanj/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 493, in __call__
    result = self.forward(*input, **kwargs)
  File "/u/nkinyanj/anaconda3/lib/python3.7/site-packages/torch/nn/parallel/data_parallel.py", line 152, in forward
    outputs = self.parallel_apply(replicas, inputs, kwargs)
  File "/u/nkinyanj/anaconda3/lib/python3.7/site-packages/torch/nn/parallel/data_parallel.py", line 162, in parallel_apply
    return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
  File "/u/nkinyanj/anaconda3/lib/python3.7/site-packages/torch/nn/parallel/parallel_apply.py", line 83, in parallel_apply
    raise output
  File "/u/nkinyanj/anaconda3/lib/python3.7/site-packages/torch/nn/parallel/parallel_apply.py", line 59, in _worker
    output = module(*input, **kwargs)
  File "/u/nkinyanj/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 493, in __call__
    result = self.forward(*input, **kwargs)
  File "/u/nkinyanj/anaconda3/lib/python3.7/site-packages/torchvision/models/densenet.py", line 122, in forward
    out = self.classifier(out)
  File "/u/nkinyanj/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 493, in __call__
    result = self.forward(*input, **kwargs)
  File "/dccstor/sgdermatology/scripts/sg_fair_Newton6/models.py", line 570, in forward
    ap = self.ap(x)
  File "/u/nkinyanj/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 493, in __call__
    result = self.forward(*input, **kwargs)
  File "/u/nkinyanj/anaconda3/lib/python3.7/site-packages/torch/nn/modules/pooling.py", line 1060, in forward
    return F.adaptive_avg_pool2d(input, self.output_size)
  File "/u/nkinyanj/anaconda3/lib/python3.7/site-packages/torch/nn/functional.py", line 788, in adaptive_avg_pool2d
    _output_size = _list_with_default(output_size, input.size())
  File "/u/nkinyanj/anaconda3/lib/python3.7/site-packages/torch/nn/modules/utils.py", line 22, in _list_with_default
    raise ValueError('Input dimension should be at least {}'.format(len(out_size) + 1))
ValueError: Input dimension should be at least 3
</code></pre>

PapRazzi: <p>Class Flatten is defined as</p>
<pre><code class="lang-auto">class Flatten(nn.Module):
    def forward(self, input):
        return input.view(input.size(0), -1)
</code></pre>

ptrblck: <p>Based on your model architecture, you should provide a 4-dimensional input as <code>[batch_size, channels, height, width]</code>.<br>
Both adaptive pooling layers will create an output of shape <code>[batch_size, channels, 1, 1]</code>.<br>
These outputs will be concatenated in their channel dimension and flattened, yielding an activation of <code>[batch_size, channels*2]</code>.<br>
The following <code>nn.BatchNorm1d</code> layer will normalize these channels and assume a sequence length of 1.<br>
Given that its number of channels is defined as <code>inp*2</code>, your model should work with an input of shape <code>[batch_size, 1920, h, w]</code>, where <code>batch_size&gt;=2</code> and <code>h, w &gt;= 1</code>, e.g.:<br>
<code>[2, 1920, 2, 2]</code>.</p>
<p>Your output however seems to show that you are passing an input of <code>[8, 1920]</code>.<br>
Could you explain your use case and the architecture a bit, as I’m not sure to understand it completely?</p>

PapRazzi: <p>Thanks <a class="mention" href="/u/ptrblck">@ptrblck</a>. I’m finetuning a pretrained densenet201 from PyTorch. The input image is RGB, image dimension is 600x450. I’m freezing the first layers and updating the weights of the other successive layers. When I replace the default pretrained densenet201 classifier with this class, I get the dimension error.</p>

gaurav8341: <p>I have same problem, the problem seems to be that densenet is flattening the input before sending it to classifier. <a class="mention" href="/u/paprazzi">@PapRazzi</a> Did you find any solution to this problem?</p>

dvdblk: <p><a class="mention" href="/u/gaurav8341">@gaurav8341</a> same issue here. Any solutions?</p>