Combining Trained Models in PyTorch

gewa24: <p>Hi all,<br>
I’m currently working on two models that train on separate (but related) types of data. I’d like to make a combined model that than take in an instance of each of the types of data, runs them through each of the models that was pre-trained individually, and then has a few feed-forward layers at the top that process the combined result of the two individual models.</p>
<p>So far, I know that I can modify forward to take in both inputs, so I just have copied the architectures of my individual models into the combined one, process them both separately by running them through the correct layers in forward(), and then combining the results as I described. What I can’t figure out how to do is used the <em>pre trained</em> models, rather than having to use the same architecture from scratch, in the combined model.<br>
I’d really appreciate any input! Thanks!</p>

ptrblck: <p>The logic you described seems to be reasonable.<br>
I’m not sure I understand the issue correctly, but it seems you are not sure, how to restore the pretrained models and use them in the new model.<br>
If that’s the case, you could instantiate both pretrained models, load the <code>state_dicts</code>, and pass them to your new model.<br>
I’ve created a small example, since it might be easier to just see the code:</p>
<pre><code class="lang-python">class MyModelA(nn.Module):
    def __init__(self):
        super(MyModelA, self).__init__()
        self.fc1 = nn.Linear(10, 2)
        
    def forward(self, x):
        x = self.fc1(x)
        return x
    

class MyModelB(nn.Module):
    def __init__(self):
        super(MyModelB, self).__init__()
        self.fc1 = nn.Linear(20, 2)
        
    def forward(self, x):
        x = self.fc1(x)
        return x


class MyEnsemble(nn.Module):
    def __init__(self, modelA, modelB):
        super(MyEnsemble, self).__init__()
        self.modelA = modelA
        self.modelB = modelB
        self.classifier = nn.Linear(4, 2)
        
    def forward(self, x1, x2):
        x1 = self.modelA(x1)
        x2 = self.modelB(x2)
        x = torch.cat((x1, x2), dim=1)
        x = self.classifier(F.relu(x))
        return x

# Create models and load state_dicts    
modelA = MyModelA()
modelB = MyModelB()
# Load state dicts
modelA.load_state_dict(torch.load(PATH))
modelB.load_state_dict(torch.load(PATH))

model = MyEnsemble(modelA, modelB)
x1, x2 = torch.randn(1, 10), torch.randn(1, 20)
output = model(x1, x2)
</code></pre>

gewa24: <p>Thanks so much! Exactly what I needed; I wasn’t aware I could use whole models in the forward (but I guess that makes sense, as they all subclass nn.Module).</p>

Arjan_Hada: <p><a class="mention" href="/u/ptrblck">@ptrblck</a>- Can you show an example of how would you do the same thing for let’s say densenet161 and resnet152? Really appreciate your help.</p>

smth: <p>the example above is as close as it gets to simply changing ModelA / ModelB with a ResNet / Densenet.</p>
<p>If you are doing a homework problem, I dont think we’re here to give you the exact solution and give you free points <img src="https://discuss.pytorch.org/images/emoji/apple/stuck_out_tongue.png?v=6" title=":stuck_out_tongue:" class="emoji" alt=":stuck_out_tongue:"></p>

Arjan_Hada: <p><a class="mention" href="/u/smth">@smth</a> - Not a homework problem, just trying to learn something completely out of my domain.</p>
<p>If I understand it correctly, if modelA is resnet152, I can change self.fc1 to self.fc1 = nn.Linear(2048, x) where 2048 is in_features for resnet152 and x = out_features one desires. We can do the same for modelB. But in self.classifier = nn.Linear(4,2) where did those number 4 and 2 come from.</p>
<p>I understand these are very trivial for a man of your stature in programming, but for nubes…I appreciate your prompt response and making sure you are not solving someone else’s homework.</p>

ptrblck: <p>The number of input features for <code>self.classifier</code> is defined by the sum of the output features of both “pre-ensemble” models.<br>
As you can see, <code>ModelA</code> and <code>ModelB</code> both output 2 values (from <code>self.fc1</code>). These features will be concatenated in <code>MyEnsemble</code> and passed to <code>self.classifier</code>, which takes 4 input features.</p>

Arjan_Hada: <p>Thanks <a class="mention" href="/u/ptrblck">@ptrblck</a> for clarification. Got it finally.</p>

himahuja: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> If I call <code>eval()</code> on a <code>MyEnsemble</code>  object, will it automatically call eval on <code>MyModelA</code> and <code>MyModelB</code>?</p>
<p>Thank you for your answer! <img src="https://discuss.pytorch.org/images/emoji/apple/100.png?v=6" title=":100:" class="emoji" alt=":100:"></p>

bhushans23: <p>Yes.<br>
eval recursively calls eval on its children nodes.<br>
Ref- <a href="https://pytorch.org/docs/stable/_modules/torch/nn/modules/module.html#Module.eval" rel="nofollow noopener">https://pytorch.org/docs/stable/_modules/torch/nn/modules/module.html#Module.eval</a></p>

Kirayue: <p><a class="mention" href="/u/ptrblck">@ptrblck</a>, for you example above. Is the weights of ModelA and ModelB is frozen or trained with classifier?</p>

ptrblck: <p>All parameters will be trained end-to-end in the example.</p>

jmlb: <p><a class="mention" href="/u/ptrblck">@ptrblck</a>, thank you for the toy example. How would you change the script if <code>modelB</code> 's input was taken from one of the layers output of <code>modelA</code>. And the ensemble output has 2 heads.In the modification below, <code>fc1</code> of <code>modelA</code> is ran twice (not efficient):</p>
<pre><code class="lang-auto">class MyModelA(nn.Module):
    def __init__(self):
        super(MyModelA, self).__init__()
        self.fc1 = nn.Linear(20, 10)
        self.fc2 = nn.Linear(10, 2)
        
    def forward(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        return x
    

class MyModelB(nn.Module):
    def __init__(self):
        super(MyModelB, self).__init__()
        self.fc1 = nn.Linear(10, 2)
        
    def forward(self, x):
        x = self.fc1(x)
        return x

class MyEnsemble(nn.Module):
    def __init__(self, modelA, modelB):
        super(MyEnsemble, self).__init__()
        self.modelA = modelA
        self.modelB = modelB
        
    def forward(self, x1, x2):
        head1 = self.modelA(x1)
        head2 = self.modelB(self.modelA.fc1(x1))
        x = torch.cat((head1, head2), dim=1)
        return x
</code></pre>
<p>Thanks!</p>

ptrblck: <p>You could return both values in the <code>forward</code> method of <code>ModelA</code>:</p>
<pre><code class="lang-python">class MyModelA(nn.Module):
    def __init__(self):
        super(MyModelA, self).__init__()
        self.fc1 = nn.Linear(20, 10)
        self.fc2 = nn.Linear(10, 2)
        
    def forward(self, x):
        x1 = self.fc1(x)
        x2 = self.fc2(x1)
        return x1, x2

...

class MyEnsemble(nn.Module):
    def __init__(self, modelA, modelB):
        super(MyEnsemble, self).__init__()
        self.modelA = modelA
        self.modelB = modelB
        
    def forward(self, x1, x2):
        head1a, head1b = self.modelA(x1)
        head2 = self.modelB(head1a)
        x = torch.cat((head1b, head2), dim=1)
        return x
</code></pre>

jmlb: <p>I see! That requires modifying the forward method of  <code>modelA</code>?<br>
Here is a few words about what I am trying to do.<br>
I have <code>ModelA</code> :<br>
<code>[input/image -&gt; encoder1 -&gt; header1 -&gt; output1]</code>.<br>
The model was trained and has pretty good accuracy.<br>
Now I want to add a 2nd task.<br>
For that purpose, I run the entire dataset through <code>ModelA</code> and for each example, save the  activation at the last layer of the encoder as numpy array file.<br>
Those numpy arrays are then used as the inputs to <code>modelB</code>  <code>[input2 -&gt; header2 -&gt; output2]</code>. Then run optimization and save the weights of <code>modelB</code>.</p>
<p>Finally, I want to combine <code>modelA</code> and <code>modelB</code> into a combo-model with 2 heads:</p>
<pre><code class="lang-auto">combo_model:  [input/image -&gt; encoder1 -&gt; header1 -&gt; output1 
                                      \-&gt; header2 -&gt; output2
</code></pre>
<p>This <code>combo_model</code> is only used at test time.<br>
(I understand I could start with the <code>combo model</code> architecture and train the <code>modelA</code> part with freezing <code>modelB</code> branch and then optimize <code>modelB</code> with freezing <code>ModelA</code>. IMHO, the approach is more cumbersome!)</p>

ptrblck: <p>You don’t need to retrain the models, but could just redefine the <code>forward</code> of <code>modelA</code> and load it’s <code>state_dict</code>.<br>
The <code>state_dict</code> contains only the parameters of the model, so you are free to manipulate the actual <code>forward</code> method as you wish, since the computation graph is created dynamically. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

jmlb: <p>OK, thank you for your answers.</p>

vainaijr: <p>how will one continue creating these ensembles, of 2 to 4 to 8 neural networks and so on?</p>

Tarlison: <p>That’s helped  a lot. Thanks!</p>

viduzz84: <p>Hi, i’m working with the exact same structure of your example above. How can i keep frozen the weights of the two models A and B (pre-trained) and train only the classifier layer of the ensemble model?</p>
<p>thanks</p>