How to alternate inference with modified model and training?

username: Alex_I
cooked: <p>I’d like to alternate inference and training in a Resnet-type model, but I need to reconfigure it during inference, and I have some questions about that.</p>
<ol>
<li>
<p>Can I unload the model from the gpu by calling model.to(‘cpu’), make a modified copy (and run it on the gpu), and then move the original back to gpu by calling model.to(‘gpu’)?  In other words, is moving the model gpu-&gt;cpu-&gt;gpu a lossless operation?  (And what happens to the parameters that were passed to the optimizer?  I don’t want to lose the optimizer state)</p>
</li>
<li>
<p>What is the best way to make a copy of an in-memory model?  I can save it and then reload a copy, but not sure if that is necessary just to copy.</p>
</li>
<li>
<p>If I want to run inference in half precision (more than 2x faster in this case), can I change the model to half and then change it back?  Is that lossless?  (ie - does the model keep a copy of everything, or does it replace weights with half-precision in place?)</p>
</li>
</ol>

username: ptrblck
cooked: <ol>
<li>
<p>I’m not sure what “lossless” means in this context. The optimizer stored references to the model parameters, so depending what your modification does, you might delete these references. Generally, you can move models between different devices (and this operation is also differentiable).</p>
</li>
<li>
<p>I would recommend to create a <code>deepcopy</code> of the <code>state_dict</code>.</p>
</li>
<li>
<p>Calling <code>half()</code> on the model might work, but could also create NaN/Inf outputs since FP16 can easily overflow. In particular during training we thus recommend to use <a href="https://pytorch.org/docs/master/amp.html">automatic mixed-precision</a>, but your inference use case might also work.<br>
All parameters will be transformed to <code>HalfTensors</code> and back to <code>FloatTensor</code> and will thus lose precision as seen in this example:</p>
</li>
</ol>
<pre><code class="lang-python">x = torch.randn(100)
y = x.half().float()
print((x - y).abs().max())
</code></pre>

username: Alex_I
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> - Thanks!  By lossless I mean a model can be round-tripped gpu-cpu-gpu and resume training as though the round trip did not take place - the weights are exactly the same, the optimizer state is the same, and nothing needs to be set up again.  I’m not sure what you mean by moving between devices is differentiable, but it sounds like the optimizer should continue working after a round trip.</p>
<p>Putting together all of your suggestions - the train/inference/train loop could look something like this in simplified code:</p>
<pre><code class="lang-auto"># different models, but with compatible weights
model = ...
inference_model = ...
optimizer = optim.AdamW(model.parameters(), ...)

while True:

    model.to(device)
    model.train()
    
    for (data, target) in loader:
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        logits = model(data)
        loss = loss_func(logits, target)
        loss.backward()
        optimizer.step()
        
    model.to('cpu')

    inference_model.load_state_dict(copy.deepcopy(model.state_dict())) # is deepcopy needed here?
    inference_model.half() # doesn't affect original model's weights
    inference_model.to(device)
    inference_model.eval()
    
    with torch.no_grad():
        for (data, target) in loader:
            data, target = data.to(device), target.to(device)
            logits = inference_model(data)

    inference_model.to('cpu') # maybe not necessary, just need to unload it, don't need a copy of the weights
</code></pre>
<p>Does that look about right?</p>
<p>Main things I’m not sure about<br>
(a) Really don’t need to do anything with the optimizer, it would “just work”?   I  now realize in typical code the optimizer is created with model.parameters() before model.to(device), so it must work across moves <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:">  Are there any implications of doing a <em>lot</em> of moves vs just one?<br>
(b) Copy state dict from one model to another like shown?</p>

username: ptrblck
cooked: <p>Thanks for the clarification.</p>
<ol>
<li>A stateless optimizer might work, but if it’s using internal buffers, such as e.g. <code>Adam</code>, you would get a device mismatch error. I would generally recommend to recreate the optimizer and load its <code>state_dict</code>. An annoyance is that you would need to push the internal states to the new device. Here is a code snippet adapted from <a href="https://github.com/pytorch/pytorch/issues/2830#issuecomment-336194949">this issue</a> (a <code>to()</code> method for optimizers might be useful here <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=9" title=":confused:" class="emoji" alt=":confused:"> ):</li>
</ol>
<pre><code class="lang-python">model = nn.Linear(1, 1).cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# Single update
out = model(torch.randn(1, 1).cuda())
out.mean().backward()
w0 = model.weight.clone()
optimizer.step()
w1 = model.weight.clone()
print((w1 - w0).abs().max())
optimizer.zero_grad()

# Move to CPU and verify update
sd = optimizer.state_dict()
model.cpu()
optimizer_cpu = torch.optim.Adam(model.parameters(), lr=1e-3)
optimizer.load_state_dict(sd)
for state in optimizer.state.values():
    for k, v in state.items():
        if isinstance(v, torch.Tensor):
            state[k] = v.cpu()
            
out = model(torch.randn(1, 1))
out.mean().backward()
w0 = model.weight.clone()
optimizer.step()
w1 = model.weight.clone()
print((w1 - w0).abs().max())
</code></pre>
<ol start="2">
<li>Yes, the <code>deepcopy</code> seems to be fine here. The <code>load_state_dict</code> method would call <code>param.copy_</code> as seen <a href="https://github.com/pytorch/pytorch/blob/488ee3790e6812fa9e940cab23e556bd3fdf330f/torch/nn/modules/module.py#L979">here</a>, so it shouldn’t be necessary to create the deepcopy of the <code>state_dict</code>, but it might not hurt. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:">
</li>
</ol>