Densenet Transfer Learning

username: balaji
cooked: <p>Hi all,<br>
I was able to implement Resnet34 using the class as follows</p>
<pre><code class="lang-python">class Resnet34(nn.Module):
    def __init__(self, num_classes = 10):
        super(Resnet34,self).__init__()
        original_model = models.resnet34(pretrained=True)
        self.features = nn.Sequential(*list(original_model.children())[:-1])
        self.classifier = nn.Sequential(nn.Linear(512, num_classes))

    def forward(self, x):
        f = self.features(x)
        f = f.view(f.size(0), -1)
		y = self.classifier(f)
        return f,y
</code></pre>
<p>But when i try similar approach for densenet.</p>
<pre><code class="lang-python">class Densenet161(nn.Module):
    def __init__(self, num_classes = 2):
        super(Densenet161,self).__init__()
        original_model = models.densenet161(pretrained=True)
        self.features = nn.Sequential(*list(original_model.children())[:-1])
        self.classifier = (nn.Linear(2208, num_classes))

    def forward(self, x):
        f = self.features(x)
        f = f.view(f.size(0), -1)
		y = self.classifier(f)
        return f,y
</code></pre>
<p>I get the following error</p>
<pre><code class="lang-auto">RuntimeError: size mismatch at /py/conda-bld/pytorch_1493676237139/work/torch/lib/THC/generic/THCTensorMathBlas.cu:243
</code></pre>
<p>I am using the same code used in Pytorch transfer learning fine tuning example.<br>
Please suggest a method to sort the issue</p>

username: smth
cooked: <p>It is best to read the model definitions first, otherwise you wont know what’s wrong.<br>
Densenet has average pooling before the classifier stage:<br>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/vision/blob/master/torchvision/models/densenet.py#L158-L159" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/vision/blob/master/torchvision/models/densenet.py#L158-L159" target="_blank" rel="nofollow noopener">pytorch/vision/blob/master/torchvision/models/densenet.py#L158-L159</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="158" style="counter-reset: li-counter 157 ;">
<li>if isinstance(m, nn.Conv2d):</li>
<li>    nn.init.kaiming_normal(m.weight.data)</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>
<p>Here’s a corrected version:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
from torch.autograd import Variable
import torchvision.models as models
import torch.nn.functional as F

class Densenet161(nn.Module):
    def __init__(self, num_classes = 2):
        super(Densenet161,self).__init__()
        original_model = models.densenet161()
        self.features = nn.Sequential(*list(original_model.children())[:-1])
        self.classifier = (nn.Linear(2208, num_classes))

    def forward(self, x):
        f = self.features(x)
        f = F.relu(f, inplace=True)
        f = F.avg_pool2d(f, kernel_size=7).view(f.size(0), -1)
        y = self.classifier(f)
        return f,y


x = Densenet161()

inp = Variable(torch.randn(4, 3, 224, 224))

out = x(inp)
</code></pre>

username: balaji
cooked: <p>Thanks <a class="mention" href="/u/smth">@smth</a> … I downloaded the densenet code from the pytorch git and I modified the model as per my need. Its working good. Thanks</p>

username: muralidharpettela
cooked: <p>Here is the correct code for Transfer learning with densenet<br>
dset_classes_number = len(class_names)<br>
model_ft = models.densenet121(pretrained=True)<br>
model_ft.classifier = nn.Linear(1024, dset_classes_number)</p>

username: 111104
cooked: <p>Thanks! Work perfectly!</p>