Loading saved Amp models into ensemble

username: Eiphodos
cooked: <p>I have trained several models with amp in FP16 separately and saved all state dicts for both models and amp.</p>
<p>Now as a continuation of that training, I would like to load those models into an ensemble, freeze all gradients and train a new model with a few final layers that learns based on the output from those models and several different inputs.</p>
<p>My current code looks something like this</p>
<pre><code class="lang-auto">modelA = ModelA()
modelB = ModelB()

modelA.load_state_dict(checkpointA['model'])
modelB.load_state_dict(checkpointB['model'])

for a in modelA.parameters():
    a.requires_grad = False
for b in modelB.parameters():
    b.requires_grad = False

ensemble = EnsembleModel(modelA, modelB)
optimizer = FusedAdam(filter(lambda p: p.requires_grad, ensemble .parameters()), lr=learning_rate)

ensemble , optimizer = amp.initialize(ensemble , optimizer, opt_level)

***perform training***
</code></pre>
<p>Now to my questions:</p>
<p>In the checkpoints for modelA and modelB I also saved the amp state dict. Should those be loaded in this case and how?<br>
Does amp.initialize support models embedded inside an ensemble?<br>
What if I want to unfreeze the lower models at a later point, do I need to reinitialize both the optimizer and amp in that case?</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="Eiphodos" data-post="1" data-topic="97338">
<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/eiphodos/48/30974_2.png" class="avatar"> Eiphodos:</div>
<blockquote>
<p>In the checkpoints for modelA and modelB I also saved the amp state dict. Should those be loaded in this case and how?</p>
</blockquote>
</aside>
<p>This wouldn’t be necessary, if you are not planning to finetune the pretrained models.</p>
<aside class="quote no-group" data-username="Eiphodos" data-post="1" data-topic="97338">
<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/eiphodos/48/30974_2.png" class="avatar"> Eiphodos:</div>
<blockquote>
<p>Does amp.initialize support models embedded inside an ensemble?</p>
</blockquote>
</aside>
<p>We recommend to use the native amp implementation via <a href="https://pytorch.org/docs/stable/amp.html"><code>torch.cuda.amp</code></a> instead of apex/amp. Since there is no <code>amp.initialize</code> method, it should just work, but let us know if you encounter any issues.</p>
<aside class="quote no-group" data-username="Eiphodos" data-post="1" data-topic="97338">
<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/eiphodos/48/30974_2.png" class="avatar"> Eiphodos:</div>
<blockquote>
<p>What if I want to unfreeze the lower models at a later point, do I need to reinitialize both the optimizer and amp in that case?</p>
</blockquote>
</aside>
<p>The amp <code>state_dict</code> stores the loss scaler values etc. which is useful to continue the training with the same setup. If you initialize a new <code>scaler</code>, the first iterations might be skipped, if the scaling factor is too high, which wouldn’t necessarily be the case using the scaling factor from the last iteration in the pretraining.</p>

username: Eiphodos
cooked: <p>I see, thanks for the help!</p>