RuntimeError: mat1 dim 1 must match mat2 dim 0 for Sequential of DINO model

Tizian_Th: <p>I’m getting the runtime error when trying to Sequentialise my linear layer with a DINO backbone from torch.hub.</p>
<pre><code class="lang-auto">import torch
import torch.distributed as dist

class LinearClassifier(torch.nn.Module):
    def __init__(self, dim, num_labels=1000):
        super(LinearClassifier, self).__init__()
        self.num_labels = num_labels
        self.linear = torch.nn.Linear(dim, num_labels)
        self.linear.weight.data.normal_(mean=0.0, std=0.01)
        self.linear.bias.data.zero_()

    def forward(self, x):
        # flatten
        x = x.view(x.size(0), -1)

        # linear layer
        return self.linear(x)


dist.init_process_group('gloo', init_method='file:///tmp/somefile', rank=0, world_size=1)
# load backbone
model = torch.hub.load('facebookresearch/dino:main', 'dino_vits8')

#Setup linear layer
linear_classifier = LinearClassifier(1536, 1000)
linear_classifier = linear_classifier.cuda()
linear_classifier = torch.nn.parallel.DistributedDataParallel(linear_classifier)
state_dict = torch.hub.load_state_dict_from_url(url="https://dl.fbaipublicfiles.com/dino/dino_deitsmall8_pretrain/dino_deitsmall8_linearweights.pth")['state_dict']
linear_classifier.load_state_dict(state_dict, strict=True)

#Sequentialise
model = torch.nn.Sequential(model,
                            linear_classifier)

x = torch.ones((1, 3, 224, 224))
out = model(x)
print("out: " + out)
</code></pre>
<p>This is what the print of the final layers looks like:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/2/32e81bb3beb656482a09c9d588a766df72ff0988.png" data-download-href="https://discuss.pytorch.org/uploads/default/32e81bb3beb656482a09c9d588a766df72ff0988" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/3/2/32e81bb3beb656482a09c9d588a766df72ff0988.png" alt="image" data-base62-sha1="7gla6QCvmYz15CqsZxRjBC1ImHm" width="437" height="499" data-dominant-color="0F305F"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1084×1240 44.7 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

ColdFrenzy: <p><a class="mention" href="/u/tizian_th">@Tizian_Th</a><br>
The output of the Dino model is a [1,384] tensor, so if you want to connect two sequential layers, the output of the previous should have the same dimension as the input of the next. If you substitute</p>
<pre><code class="lang-auto">linear_classifier = LinearClassifier(384, 1000)
</code></pre>
<p>It should work</p>

Tizian_Th: <p>I tried that first but then I got the error:<br>
RuntimeError: Error(s) in loading state_dict for DistributedDataParallel: size mismatch for module.linear.weight: copying a param with shape torch.Size([1000, 1536]) from checkpoint, the shape in current model is torch.Size([1000, 384]).</p>

ColdFrenzy: <p>ok, clear.<br>
The pre-trained layer you are using has dimensions [1000, 1536], the ideas that come to mind are 3:</p>
<ol>
<li>if you really need that pre-trained layer you could add an intermediate layer in the sequential in order to re-expand the output as:</li>
</ol>
<pre><code class="lang-auto">model = torch.nn.Sequential(model,torch.nn.Linear(384,1536),
                            linear_classifier)

</code></pre>
<p>although, in this case you would still have to re-train the network a bit to tune the parameters of the new layer.</p>
<ol start="2">
<li>As in my previous answer, you can remove these two lines</li>
</ol>
<pre><code class="lang-auto">state_dict = torch.hub.load_state_dict_from_url(url="https://dl.fbaipublicfiles.com/dino/dino_deitsmall8_pretrain/dino_deitsmall8_linearweights.pth")['state_dict']
linear_classifier.load_state_dict(state_dict, strict=True)
</code></pre>
<p>and just use a new layer which you have to train</p>
<pre><code class="lang-auto">linear_classifier = LinearClassifier(384, 1000)
</code></pre>
<ol start="3">
<li>You can create a new dictionary with just part of the weights that you have:</li>
</ol>
<pre><code class="lang-auto">new_state_dict = {}
new_state_dict["module.linear.bias"] = state_dict["module.linear.bias"][0:384]
new_state_dict["module.linear.weight"] = state_dict["module.linear.weight"][0:384]
</code></pre>
<p>where you can take any set of 384 values (e.g. 1:385, 2:386 etc.). Anyway I don’t see the point in doing this, since you should probably fine-tune these weights again</p>

Tizian_Th: <p>Thanks for your answer this helps a lot. I am trying to recreate Facebooks results as closely as possible as I am benchmarking DINO as a part of my Bachelor Thesis. Facebook seems to reshape their output like this:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/a/8ad6cd4ccf23b33073a8484d726a96ad23389408.png" data-download-href="https://discuss.pytorch.org/uploads/default/8ad6cd4ccf23b33073a8484d726a96ad23389408" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/a/8ad6cd4ccf23b33073a8484d726a96ad23389408_2_690x186.png" alt="image" data-base62-sha1="jOe6k1JhCoOo3izpYZ8Fvwd4Kus" width="690" height="186" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/a/8ad6cd4ccf23b33073a8484d726a96ad23389408_2_690x186.png, https://discuss.pytorch.org/uploads/default/optimized/3X/8/a/8ad6cd4ccf23b33073a8484d726a96ad23389408_2_1035x279.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/8/a/8ad6cd4ccf23b33073a8484d726a96ad23389408_2_1380x372.png 2x" data-dominant-color="FAFAFB"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1816×490 46.7 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>What way would you recommend to replicate the original model as close as possible?</p>