Fine tuning part of a pre trained model which fed a new classifier

username: catt_ale
cooked: <p>Hello to everyone!</p>
<p>I have a pre-trained model made by myself, now I pick up some layers of the base model and use them to feed a classifier. Now I want to train the classifier and fine-tune the layers of my pre-trained net. My question is: is it correct how I would like to do here in the code below? <strong>I am not sure where there is the comment</strong></p>
<blockquote>
<p>#???</p>
</blockquote>
<p>Someone can explain if, where, and why I am going wrong?<br>
Thank you so much!</p>
<pre data-code-wrap="python"><code class="lang-python">class Mybasemodel(nn.Module):
    def __init__(self, modelA, modelB):
        self.model = nn.Sequential(
            .
            .
            .
        )
    def forward(self, x1, x2):
        x = self.model(x)

        return x


class Classifier(nn.Module):
    def __init__(self, modelA, modelB):
        self.model = nn.Sequential(
            .
            .
            .
        )
    def forward(self, x1, x2):
        x = self.model(x)

        return x

basemodel = Mybasemodel()
basemodel.load_state_dict(torch.load('...'))
basemodel.train()

calssifier = Classifier()
calssifier.train()

loss_basemodel_fn = torch.nn.L1Loss()
loss_classifier_fn = torch.nn.L1Loss()

opt_basemodel = torch.optim.(basemodel.parameters())
opt_classifier = torch.optim.(calssifier.parameters())

first_layer = 0
N_layer = 3

i = 0
for param in model.basemodel():
    i += 1
    if i &gt; N_layer:
        param.requires_grad = False


for i, data in enumerate(dataset):

    input = data['input']
    label_basemodel = data['label_basemodel']
    label_classifier = data['label_classifier']

    output_base_model = basemodel.model(input)

    feature_map = basemodel.model[first_layer:N_layer](input)
    pred = calssifier(feature_map)

    opt_basemodel.zero_grad()
    opt_classifier.zero_grad()

    loss_classifier = loss_classifier_fn(pred, label_basemodel)
    #???
    loss_basemodel = loss_basemodel_fn(output_base_model, label_classifier) + loss_classifier

    loss_basemodel.backward()
    loss_classifier.backward()

    opt_basemodel.step()
    opt_classifier.step()

</code></pre>

username: ptrblck
cooked: <p>It depends on your use case and how you would like to fine tune the <code>basemodel</code>.<br>
Currently both losses would create gradients in the parameters or <code>basemodel</code>, since these parameters were used to compute the losses. If you only want to create the gradients of <code>loss_basemodel</code> w.r.t to the parameters of <code>basemodel</code>, you would have to <code>.detach()</code> <code>output_base_model</code> before feeding it to <code>classifier</code>.</p>
<p>Note that there are minor issues (such as using <code>for param in model.basemodel()</code> instead of <code>for param in basemodel.parameters()</code>) but I assume these are copy-paste errors.<br>
Also, if you are iterating the parameters, <code>i</code> will not reflect the layers, since some modules have more than a single parameter (often <code>weight</code> and <code>bias</code>).</p>

username: catt_ale
cooked: <p>Hello, thank you for your answer.</p>
<p>I don’t understand why would I detach <code>output_base_model</code> before feeding it to the classifier? This sounds new to me.</p>
<p>To give you some additional information, my base model is a cycle GAN, I am using some layer of the generator to reuse learned features to fed a classifier. Train the two networks separately is not a problem, I guess. With the simplified example above I wanted to show how I would optimize base-model parameters w.r.t. the classification error. The optimal for me would be able both to:</p>
<ol>
<li>Optimize the base model w.r.t his original loss function combined with the loss of the classifier</li>
<li>Optimize the base model w.r.t. the loss of the classifier.</li>
</ol>