Extract weights of specific module from pretrained model file

username: aman_goyal
cooked: <p>Greetings,</p>
<p>This is Aman Goyal. I am currently pursuing research in CMU and MSU.<br>
I actually have a pre-trained model file which basically has 3 modules: backbone, neck and head. It is basically a <a href="https://github.com/SysCV/qdtrack" rel="noopener nofollow ugc">deep learning based tracking method</a>.<br>
Now what I need to do is that I just want it’s backbone weights which is basically a Resnet-50 trained on BDD100K detection dataset.<br>
Now is there any possibility of extraction of only backbone weights from the <a href="https://drive.google.com/file/d/1YNAQgd8rMqqEG-fRj3VWlO4G5kdwJbxz/view?usp=sharing" rel="noopener nofollow ugc">.pth model file</a> which basically contains all modules trained in end-to-end fashion.</p>

username: JuanFMontesinos
cooked: <p>The .pth file contains a sorted dictionary. Keys are sorted the same way you called them while instantiating the model.<br>
This basically means if you call model.backbone.state_dict().keys() you should get which ones correspond to the backbone.<br>
Then you just get them from the pth file.</p>

username: sagorsarker
cooked: <p>Hello <a class="mention" href="/u/aman_goyal">@aman_goyal</a> ,<br>
If I understand correctly, you can extract provided model backbone weights by the following codes:</p>
<pre><code class="lang-auto">import torch
model = torch.load('qdtrack-frcnn_r50_fpn_12e_bdd100k-13328aed.pth')
new_state_dict = {}
for name, weight in model['state_dict'].items():
    if 'backbone' in name:
        new_state_dict[name] = weight
</code></pre>
<p>This new state dict contains the end-to-end trained weights.  you can create a new model with this state dict and add a new neck or head I guess.</p>

username: aman_goyal
cooked: <p>Thanks a lot for your prompt response. Can I create a new .pth file consisting of only backbone weights ?</p>

username: sagorsarker
cooked: <p>sure you can.<br>
I think you can remove the other layer and keep only the backbone network.<br>
Then save it with your desired name.</p>
<pre><code class="lang-auto">new_model = {
    'meta': model['meta'],
    'state_dict': new_state_dict
}
torch.save(new_model, "backbone.pth")
</code></pre>
<p>Hope this will help.</p>

username: aman_goyal
cooked: <p>Could you mention the code snippet which I would be required to use for saving the backbone weights as .pth file</p>

username: sagorsarker
cooked: <p>Hello <a class="mention" href="/u/aman_goyal">@aman_goyal</a> ,<br>
please check my previous reply.<br>
I have edited with a code snippet.</p>

username: aman_goyal
cooked: <p><a class="mention" href="/u/sagorsarker">@sagorsarker</a>  Hey, so I am slightly confused with the 2 code snippets now. So how should they be combined to get backbones.pth ? Could you please combine them in a single comment for more clarity.</p>
<p>Thanks</p>

username: sagorsarker
cooked: <p>Hi <a class="mention" href="/u/aman_goyal">@aman_goyal</a> ,<br>
Here is the full code snippet about how I tried.</p>
<pre><code class="lang-auto">import torch
from collections import OrderedDict

# load the model
model = torch.load('qdtrack-frcnn_r50_fpn_12e_bdd100k-13328aed.pth')

# separate backbone layers
new_state_dict = {}
for name, param in model['state_dict'].items():
    if 'backbone' in name:
        new_state_dict[name] = param

new_state_dict = OrderedDict(new_state_dict)

# construct new model with backbone layers
new_model = {
    'meta': model['meta'],
    'state_dict': new_state_dict,
}

# save the new model which has only backbone layers
torch.save(new_model, 'backbone.pth')
</code></pre>
<p>If this works let me know.</p>

username: aman_goyal
cooked: <p>I did get ‘backbone.pth’ out of this but is there any way to test it to just verify whether it is correct or not?</p>