Fine-tuning RetinaNet for user-defined num of classes

Adrian_Contreras_Gar: <p>Hi!</p>
<p>I’m not quite experienced on modifying pretrained models’ architecture, but this time for my application I need it. I want to change the classification head of the retinanet_resnet50 model in order to adapt for a dataset with 6 classes. Any idea of how can i code that?</p>

KFrank: <p>Hi Adrian!</p>
<aside class="quote no-group" data-username="Adrian_Contreras_Gar" data-post="1" data-topic="202126" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/adrian_contreras_gar/48/62351_2.png" class="avatar"> Adrian_Contreras_Gar:</div>
<blockquote>
<p>I want to change the classification head of the retinanet_resnet50 model in order to adapt for a dataset with 6 classes.</p>
</blockquote>
</aside>
<p>I’ve never used pytorch’s <code>RetinaNet</code>, but it appears that you can instantiate one<br>
with a pre-trained <code>ResNet50</code> backbone with a user-specified number of classes.<br>
Doing this won’t load pre-trained weights for the classification head (which makes<br>
sense), so you’ll have to train those weights from scratch.</p>
<p>(It would probably make sense to train the classification-head weights some while<br>
holding the backbone weights fixed so that you don’t adapt the pre-trained backbone<br>
weights to the untrained head weights, after which you could fine tune the whole<br>
model for a while.)</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'2.3.0'
&gt;&gt;&gt; import torchvision
&gt;&gt;&gt; torchvision.__version__
'0.18.0'
&gt;&gt;&gt; resnet6 = torchvision.models.detection.retinanet_resnet50_fpn (weights_backbone = 'DEFAULT', num_classes = 6)
Downloading: "https://download.pytorch.org/models/resnet50-11ad3fa6.pth" to &lt;path_to_pytorch_cache&gt;\torch\hub\checkpoints\resnet50-11ad3fa6.pth
100.0%
&gt;&gt;&gt; resnet6.head.classification_head.num_classes
6
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

Adrian_Contreras_Gar: <p>Thanks for the response!</p>
<p>I tried this but when I execute the training loop, loss doesn’t decrease, which wasn’t happening when I used the pretrained weights. Is there any chance to keep these pretrained weights while adapting the model for the desired number of classes?</p>

KFrank: <p>Hi Adrian!</p>
<aside class="quote no-group" data-username="Adrian_Contreras_Gar" data-post="3" data-topic="202126" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/adrian_contreras_gar/48/62351_2.png" class="avatar"> Adrian_Contreras_Gar:</div>
<blockquote>
<p>I tried this</p>
</blockquote>
</aside>
<p>What specifically did you try?</p>
<aside class="quote no-group">
<blockquote>
<p>when I execute the training loop, loss doesn’t decrease</p>
</blockquote>
</aside>
<p>Are you trying to train the whole backbone or just the classification-head parameters?<br>
Does the loss change at all?  Do the parameters you think you are training change at<br>
all?  How long did you train for?  What optimizer and parameters are you using?</p>
<p>Can you set things up to freeze all but the classification-head parameters and then<br>
try to vigorously overfit the classification head to a smallish subset of your training<br>
data?  If you do that – training a lot – does your loss start to come down?</p>
<p>What loss function are you using?  Since you’re trying to change the number of<br>
predicted classes, does the loss have a specific class-prediction term that you can<br>
piece out and look at individually?</p>
<aside class="quote no-group">
<blockquote>
<p>which wasn’t happening when I used the pretrained weights.</p>
</blockquote>
</aside>
<p>How, specifically, did you use the pretrained weights?  In this case, are you talking<br>
about the model that predicts 91 different classes?</p>
<p>The two models for two different numbers of classes are, of course, slightly different.<br>
As a consistency test, can you randomly reinitialize the classification-head weights<br>
of the 91-class model and see whether the “loss decreases” when you train with<br>
whatever scheme you were using when you “used the pretrained weights?”</p>
<aside class="quote no-group">
<blockquote>
<p>Is there any chance to keep these pretrained weights while adapting the model for the desired number of classes?</p>
</blockquote>
</aside>
<p>Possibly.  But note that the shape of the weight tensor in the final layer of the<br>
classification head differs between the pretrained and adapted model.  So you<br>
can’t just blindly reuse those particular weights.</p>
<p>Best.</p>
<p>K. Frank</p>