Use an existing/pre-trained model as a block of a new model

pylearn: <p>Suppose I have a model, <code>modelA</code>. Can I use it as a block/layer in a new model, <code>modelB</code>, by directly assigning <code>modelA</code> as a layer in the <code>__init__()</code> of <code>ModelB</code>? What are the possible side effect? For example,</p>
<pre><code class="lang-auto">modelA = nn.Conv2d(20, 20, 5)

class ModelB(nn.Module):
    def __init__(self, modelA):
        super().__init__()
        self.conv = nn.Conv2d(1, 20, 5)
        self.convA = modelA

    def forward(self, x):
        # ...

modelB = ModelB(modelA)
</code></pre>
<p>I am concerned since the <code>modelB</code> defined above is not conventional,</p>
<ul>
<li><code>modelA</code> is initialized outside <code>__init__()</code> of <code>ModelB</code></li>
</ul>
<p>I did some tests, no error was found. However, the examples I used are simple and I am not sure if I miss anything.</p>
<p>Is there any difference if <code>modelA</code> is initialized inside <code>__init__()</code> of <code>ModelB</code>?</p>
<pre><code class="lang-auto">class ModelB(nn.Module):
    def __init__(self, modelA):
        super().__init__()
        self.conv = nn.Conv2d(1, 20, 5)
        self.convA = nn.Conv2d(20, 20, 5)
</code></pre>
<p>Further, what if <code>modelA</code> is initialized outside and some training is done before it is used in <code>modelB</code>?</p>

ptrblck: <p>Your approach will work and there is no difference where the submodules were initialized. Creating <code>modelA</code> outside of <code>modelB</code> allows you to use it separately (<code>modelB</code> holds a reference to <code>modelA</code> so will see the changes). This usage would break e.g. <code>DDP</code>, but I don’t see any issues if a) you are nor using <code>modelA</code> outside of <code>modelB</code> or if you are using a plain <code>nn.Module</code> without <code>DDP</code> etc.</p>

pylearn: <p>Thank you for your explanation. I don’t use <code>DDP</code>.</p>
<p>Just to double check the “reference” you mentioned. To my understanding, “reference” means <code>modelA</code> and <code>modelB.modelA</code> are the same object. Hence, if I trained <code>modelA</code> outside or tuned the parameters, <code>modelB.modelA</code> will change accordingly, and vice versa.</p>
<p>You also mentioned “using <code>modelA</code> outside of <code>modelB</code>”, may I ask what does it mean? I thought <code>modelA</code> and <code>modelB.modelA</code> are the same. Hence, using <code>modelA</code> is the same as using <code>modelB.modelA</code>?</p>

ptrblck: <aside class="quote no-group" data-username="pylearn" data-post="3" data-topic="187551">
<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/pylearn/48/54463_2.png" class="avatar"> pylearn:</div>
<blockquote>
<p>To my understanding, “reference” means <code>modelA</code> and <code>modelB.modelA</code> are the same object. Hence, if I trained <code>modelA</code> outside or tuned the parameters, <code>modelB.modelA</code> will change accordingly, and vice versa.</p>
</blockquote>
</aside>
<p>Yes, exactly. This can also be verified using your example:</p>
<pre data-code-wrap="python"><code class="lang-python">modelA = nn.Conv2d(20, 20, 5)

class ModelB(nn.Module):
    def __init__(self, modelA):
        super().__init__()
        self.conv = nn.Conv2d(1, 20, 5)
        self.convA = modelA

    def forward(self, x):
        x = self.conv(x)
        x = self.convA(x)
        return x

modelB = ModelB(modelA)
optimizer = torch.optim.Adam(modelB.parameters(), lr=1.)

for _ in range(3):
    optimizer.zero_grad()
    out = modelB(torch.randn(1, 1, 24, 24))
    loss = out.mean()
    loss.backward()
    optimizer.step()
    print("internal: {}".format(modelB.convA.weight.abs().sum()))
    print("external: {}".format(modelA.weight.abs().sum()))
    
# internal: 10000.4697265625
# external: 10000.4697265625
# internal: 12015.4755859375
# external: 12015.4755859375
# internal: 16073.640625
# external: 16073.640625
</code></pre>
<p>As you can see the parameters of <code>modelA</code> are updated for the “external” and “internal” <code>modelA</code> as both are referencing the same object.</p>
<aside class="quote no-group" data-username="pylearn" data-post="3" data-topic="187551">
<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/pylearn/48/54463_2.png" class="avatar"> pylearn:</div>
<blockquote>
<p>You also mentioned “using <code>modelA</code> outside of <code>modelB</code>”, may I ask what does it mean?</p>
</blockquote>
</aside>
<p>By that I meant the direct training of <code>modelA</code> and the training of <code>modelA</code> via <code>modelB</code>.</p>
<aside class="quote no-group" data-username="pylearn" data-post="3" data-topic="187551">
<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/pylearn/48/54463_2.png" class="avatar"> pylearn:</div>
<blockquote>
<p>I thought <code>modelA</code> and <code>modelB.modelA</code> are the same. Hence, using <code>modelA</code> is the same as using <code>modelB.modelA</code>?</p>
</blockquote>
</aside>
<p>Yes, but the issue I’m seeing is when more complicated workflows, such as DDP, are used.<br>
I.e. in case you are wrapping <code>modelB</code> into <code>DistributedDataParallel</code> the initialization will send the initial <code>state_dict</code> (including the state of <code>modelA</code>) to all ranks to make sure the model is exactly the same. In each training step <code>DDP</code> will then synchronize the gradients for you so that the <code>optimizer.step()</code> calls creates the same (updated) models on each rank again.<br>
If you would train <code>modelA</code> directly on any rank you could easily diverge the training, which is why I wanted to add the warning. It could still work if you are careful, but I’m also seeing an easy way to diverge the training.<br>
The same applies for quantization etc. as I’m not sure how internals work there.</p>