Combining a loaded trained model with a new model

username: rim_zayani
cooked: <p>I trained a classifier and saved the model using:</p>
<pre><code class="lang-auto">torch.save(model, "/home/zaianir/Documents/code/tuto/classif/MNIST_model.pth")
</code></pre>
<p>I’m trying to train a new classifier on top of the pretrained saved model without its last layer.<br>
I want to train only the parameters of the new added layers (I don’t want to update the saved parameters).</p>
<p>Here is some of my code:</p>
<pre><code class="lang-auto">class VAE2(nn.Module):
    def __init__(self):
        super(VAE2, self).__init__()
        self.fc3=nn.Linear(50,20)
        self.fc4=nn.Linear(20,10)
        
    def forward(self, x):    
        x=F.relu(self.fc3(x))
        x=F.relu(self.fc4(x))
        x=F.log_softmax(x)
        return x
    
    
class VAE(nn.Module):
    def __init__(self, VAE1, VAE2):
        super(VAE, self).__init__()
        self.VAE1=VAE1
        self.VAE2=VAE2
        
    def forward(self):
        x=self.VAE1
        x=self.VAE2(x)
        return x

def loss_function(inp, target):
    l=F.nll_loss(inp, target)
    return l


def train(train_dl, model, epoch_nb,lr1):   
    optimizer=torch.optim.Adam(model.parameters(), lr=lr1)
    train_loss1=[]
    for epoch in range(1,epoch_nb):
            model.train()
            train_loss=0.0
            for idx, (data, label) in enumerate(train_dl):
                data, label= data.to(device), label.to(device)
                out=model(data)
                loss=loss_function(out,data)
                train_loss+=loss.item()
                model.zero_grad()
                loss.backward()
                optimizer.step()     
            av_loss= train_loss / len(train_dl.dataset)   
            print('Epoch: {} Average loss: {:.4f}'.format(
          epoch, av_loss))


train_loader = torch.utils.data.DataLoader(
        datasets.MNIST('../data', train=True, download=True,    
                       transform=transforms.Compose([
                           transforms.ToTensor(),
                           transforms.Normalize((0.1307,), (0.3081,))
                       ])),
        batch_size=128, shuffle=True, **kwargs)
VAE1 = torch.load("/home/zaianir/Documents/code/tuto/classif/MNIST_model.pth")
VAE2_model=VAE2().to(device)
model=VAE(VAE1, VAE2_model)
epoch=30
learning_rate=0.001
train(train_loader, model, epoch, learning_rate)
</code></pre>
<p>Here are the different layers of my loaded model (VAE1):<br>
<img src="https://discuss.pytorch.org/uploads/default/original/2X/4/47a91b24dbf1480cf95d7bf385118b7d585e5ded.png" alt="LAYERS" data-base62-sha1="adWcV6cLVXOwIrxNojdkedaKQRT" width="462" height="131"></p>
<p>I’m new to pytorch and don’t know how to proceed. Is my approach correct?<br>
Thank you for your help.</p>

username: ptrblck
cooked: <p>There are some issues in your code:</p>
<ul>
<li>in <code>VAE.forward</code> you are not passing <code>x</code> to <code>self.VAE1</code>, so that <code>x</code> will be the submodule not its output</li>
<li>I would suggest to save and load the <code>state_dict</code> instead of the complete model as described in the <a href="https://pytorch.org/docs/stable/notes/serialization.html" rel="nofollow noopener">Serialization docs</a>
</li>
<li>if you would like to remove the last layer from <code>VAE1</code>, you could e.g. replace it with an <code>nn.Identity</code> layer</li>
<li>to freeze the base model’s parameters, set their <code>requires_grad</code> flag to <code>False</code> as described in the <a href="https://pytorch.org/tutorials/beginner/finetuning_torchvision_models_tutorial.html#set-model-parameters-requires-grad-attribute" rel="nofollow noopener">Finetuning tutorial</a>
</li>
</ul>

username: winterknight
cooked: <p>I have trained <code>T5ForConditionalGeneration.from_pretrained(MODEL_NAME, return_dict=True)</code> for Summarization task and another for QA task.</p>
<p>I am trying to load the models so I can use their trained encoder and decoders for a different task. The <code>summarizer_model</code> and the <code>qa_model</code> are loaded from model checkpoints. But it throws an error… When using T5ForConditionalGeneration.from_pretrained() the training starts… Is there something wrong with the way I’m using model checkpoints to train further on ?</p>
<p>As far as I can tell the issue is that I need to <code>return_dict=True</code> for the model I’ve loaded from the checkpoint.</p>
<p>Is there a way to return_dict=True for a saved model checkpoint ?</p>
<pre><code class="lang-auto">import torch.nn as nn
import torch.nn.functional as F
from collections import OrderedDict

alpha = 0.5
class CrossAttentionSummarizer(pl.LightningModule):
 def __init__(self):
       super(CrossAttentionSummarizer, self).__init__()
       self.summarizer_model = summarizer_model
       self.qa_encoder = qa_model
       self.multihead_attn = nn.MultiheadAttention(embed_dim=768, num_heads=4, batch_first=True)
       self.linear1 = nn.Linear(1024*768, 512)
       self.linear2 = nn.Linear(512, 2, bias=False)
       self.ce_loss = nn.CrossEntropyLoss()

 def forward(self, question_passage_input_ids, question_passage_attention_mask, question_labels, input_ids, attention_mask, decoder_attention_mask, labels=None):
   summarizer_output = self.summarizer_model(
       input_ids,
       attention_mask=attention_mask,
       labels=labels,
       decoder_attention_mask=decoder_attention_mask
   )

   qa_output = self.qa_encoder(
       question_passage_input_ids, 
       question_passage_attention_mask,
       question_labels
   )

   

   decoder_output = summarizer_output[3]
   encoder_output = qa_output[2]

   multi_attn_output, multi_attn_output_weights = self.multihead_attn(decoder_output, encoder_output, encoder_output)
   lin_output = self.linear1(multi_attn_output.reshape(-1, 1024*768))
   cls_outputs = self.linear2(lin_output)
   cls_outputs = nn.functional.softmax(cls_outputs, dim=1)
   cls_preds = torch.argmax(cls_outputs, dim=1)
   cls_pred_loss = self.ce_loss(cls_outputs, question_labels.type(torch.int64).squeeze(dim=1))
   return summarizer_output.loss, summarizer_output.logits, cls_pred_loss, cls_preds
</code></pre>

username: winterknight
cooked: <p>I got it!!</p>
<p>Have to save the model using <code>save_pretrained()</code>  in order to be able to load with the <code>return_dict=true</code> option.</p>
<p><img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji only-emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>