Problems loading model trained on one GPU onto another GPU with .to(device)

username: S_Ellis
cooked: <p>I have a model class myModel, which I pre-train on device cuda:1 and then save to file modelFile. I don’t understand the behaviour when trying to load this model onto another, say cuda:0. Can someone help me understand what is going on behind the scenes when one has the following:</p>
<pre><code class="lang-auto">model = myModel()
model.load_state_dict(torch.load(modelFile))
model = model.eval().to("cuda:0")
</code></pre>
<p>and why this gives different results to:</p>
<pre><code class="lang-auto">model = myModel().to("cuda:0")
model.load_state_dict(torch.load(modelFile))
model = model.eval()
</code></pre>
<p>What does it mean for a saved model to be associated with a device, and why does setting model to a different device to the saved model cause problems? Should we always load the model first and then push to a new device?</p>
<p>EDIT: I should also say that I find it slightly disturbing that such a small, inconspicuous difference in the code causes huge implications for the model output. Would it be possible to flag a warning when the model and state_dicts have different devices?</p>

username: albanD
cooked: <p>Hi,</p>
<p>Saving Tensors does save their device as it is a property of the object.<br>
You can use the <code>map_location</code> argument from <a href="https://pytorch.org/docs/stable/torch.html#torch.load">torch.load</a> to avoid any issue if the device it was saved for is not available anymore. In this case, you want to use <code>"cuda:0"</code> I guess for the new location).</p>

username: S_Ellis
cooked: <p>Thank you for replying. I understand that what you have said but I would really appreciate understanding what’s happening in the first piece of code that I wrote.<br>
The <code>model</code> is put onto cuda:0, then it loads a state dict which contains a tensor mapped to cuda:1. So what actually happens in <code>model</code>? What values will it contain after loading the state dict? If a state dict has to be loaded into a model to be useful, what is the use of it having its own device property? And would it be feasible to have a warning when the model and state dict devices do not match?</p>

username: albanD
cooked: <p>Hi,</p>
<p>When you load a state dict, you load it into the current parameters of your network. So if the network is on cuda:0, the resulting net will be on cuda:0 as well.<br>
The state dict contains Tensors and Tensors always have a device. That’s the only reason why it maintains device information.</p>
<p>In general, it should not be a problem to have them on different devices as whatever you asked for with <code>.to()</code> on the model is what counts. not what was in the state dict.</p>

username: S_Ellis
cooked: <p>Thank you for the reply again, but it doesn’t match what I’ve seen in practice. When I run the second piece of code in the original question and run the <code>model</code> net, I get results that I do not expect, implying that the code failed to move the state dict correctly to cuda:0. Is this expected behaviour, and if not should I try to make a minimum example and flag an issue?</p>

username: albanD
cooked: <p>What do you mean by “I get results that I do not expect” ? Can you be more clear about what you observe?</p>

username: S_Ellis
cooked: <p>Apologies for being vague. I had trained an autoencoder model on cuda:1, and then ran the equivalent of the second code section above. The reconstructed image from the autoencoder was not correct, it looked as if the wrong batchnorm were being applied. When I changed the code to the first section it worked fine.</p>
<p>I’m sorry, I don’t have a tidy minimum example of the issue yet, but if you think it sounds suspicious I would be happy to try to make one.</p>

username: albanD
cooked: <blockquote>
<p>When I changed the code to the first section it worked fine.</p>
</blockquote>
<p>This is quite surprising indeed. They should give the same result.<br>
I would check first that this is not about the <code>.eval()</code> mode that is set properly. It has a large impact on the result for models with batchnorm layers.</p>

username: S_Ellis
cooked: <p>I have made a minimum code (without batchnorms) that seems to show the problem, sorry that it’s a bit long:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import matplotlib.pyplot as plt

#%% define the model
class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        self.main = nn.Sequential(
            nn.ConvTranspose2d(     100, 2 * 8, 4, 1, 0, bias=False),
            nn.ReLU(True),
            nn.ConvTranspose2d(2 * 8, 2 * 4, 4, 2, 1, bias=False),
            nn.ReLU(True),
            nn.ConvTranspose2d(2 * 4, 2 * 2, 4, 2, 1, bias=False),
            nn.ReLU(True),
            nn.ConvTranspose2d(2 * 2,     2, 4, 2, 1, bias=False),
            nn.ReLU(True),
            nn.ConvTranspose2d(    2,      1, 4, 2, 1, bias=False),
            nn.Tanh()
        )

    def forward(self, input):
        output = self.main(input)
        return output

# make a version of the model on cuda:1 and save the model to file
netG = Generator().to('cuda:1')
torch.save(netG.state_dict(), 'netG.pt')
netG = netG.eval()

# load the model twice onto cuda:0, with .to(device) both before and after loading the state dict
netG_1 = Generator().to('cuda:0')
netG_1.load_state_dict(torch.load('netG.pt'))
netG_1 = netG_1.eval()

netG_2 = Generator()
netG_2.load_state_dict(torch.load('netG.pt'))
netG_2 = netG_2.eval().to('cuda:0')


#%% apply the three versions
z = torch.randn(1,100,1,1) # generator input

with torch.no_grad():
    
    output = netG(z.to('cuda:1'))
    output1 = netG_1(z.to('cuda:0'))
    output2 = netG_2(z.to('cuda:0'))

#%% display output
plt.figure()
plt.subplot(1,3,1)
plt.imshow(torch.squeeze(output.cpu()))
plt.subplot(1,3,2)
plt.imshow(torch.squeeze(output1.cpu()))
plt.subplot(1,3,3)
plt.imshow(torch.squeeze(output2.cpu()))
</code></pre>
<p>This is the figure I get when I run the code:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/a/1a2ae35585cb94358e6e7cd01976595620c37457.png" data-download-href="https://discuss.pytorch.org/uploads/default/1a2ae35585cb94358e6e7cd01976595620c37457" title="outputs"><img src="https://discuss.pytorch.org/uploads/default/original/3X/1/a/1a2ae35585cb94358e6e7cd01976595620c37457.png" alt="outputs" data-base62-sha1="3JuiGzKLknpo270SXPW0JI0tHbF" width="690" height="182" data-dominant-color="BABABA"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">outputs</span><span class="informations">1855×490 20.9 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>We can see that the middle input, with the model that was sent to cuda:0 <em><strong>before</strong></em> loading the state dict has a different output.</p>

username: albanD
cooked: <p>I don’t have a machine with 2 GPUs readily available here.<br>
While I find one, could you check if changing <code>netG_2 = netG_2.eval().to('cuda:0')</code> to <code>netG_2 = netG_2.to('cuda:0').eval()</code> change anything?</p>

username: S_Ellis
cooked: <p>I’m training models on both GPUs so can’t perform your test at the moment. I will try to do it asap and let you know the result.</p>

username: S_Ellis
cooked: <p>The results are identical in either case (making sure that the torch seed is the same for both runs too).</p>

username: albanD
cooked: <p>Ok, I’ll get a machine and check that.</p>

username: albanD
cooked: <p>Hi,</p>
<p>Sorry it took a bit of time to get the machine.<br>
I’m afraid I cannot repro with the modifies script below: I don’t have a screen on the other machine so I compare the outputs with:</p>
<pre><code class="lang-python">print(output.min(), output.mean(), output.max())
print(output1.min(), output1.mean(), output1.max())
print(output2.min(), output2.mean(), output2.max())
print((output - output1.to("cuda:1")).abs().max())
print((output - output2.to("cuda:1")).abs().max())
</code></pre>
<p>And they all seems to be the same.</p>
<p>What does these prints return for you?</p>

username: S_Ellis
cooked: <p>If I add those lines to my example and run, I get the following output:</p>
<pre><code class="lang-auto">tensor(-0.0064, device='cuda:1') tensor(-0.0003, device='cuda:1') tensor(0.0064, device='cuda:1')
tensor(-0.0040, device='cuda:0') tensor(0.0005, device='cuda:0') tensor(0.0056, device='cuda:0')
tensor(-0.0064, device='cuda:0') tensor(-0.0003, device='cuda:0') tensor(0.0064, device='cuda:0')
tensor(2.0009, device='cuda:1')
tensor(3.6893e+19, device='cuda:1')
</code></pre>
<p>So, they are not the same.</p>
<p>Even more curiously, I find that if I re-run the last two lines you sent without killing the kernel, I get:</p>
<pre><code class="lang-auto">tensor(3.6893e+19, device='cuda:1')
tensor(3.6893e+19, device='cuda:1')
</code></pre>
<p>So, for some reason re-running I get a different answer for <code>(output - output1.to("cuda:1")).abs().max()</code>.</p>
<p>Not sure where we go from here, please do let me know what other tests I should run, or if you need more information.</p>

username: albanD
cooked: <p>I ran my test with a nightly build on conda and that worked fine.<br>
Can you try a nightly build?</p>

username: S_Ellis
cooked: <p>I’ll try. Just to check because I’ve never installed a nightly build before, is this the way to do so?</p>
<p><code>conda install -c pytorch pytorch-nightly</code></p>

username: albanD
cooked: <p>I never know <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=9" title=":smiley:" class="emoji" alt=":smiley:"> I usually check here: <a href="https://pytorch.org/get-started/locally/">https://pytorch.org/get-started/locally/</a><br>
Looks like the <code>-c</code> is not at the right place <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=9" title=":smiley:" class="emoji" alt=":smiley:"></p>

username: ptrblck
cooked: <p>I also rerun the code and got:</p>
<pre><code class="lang-python">tensor(-0.0106, device='cuda:1') tensor(7.5541e-05, device='cuda:1') tensor(0.0082, device='cuda:1')
tensor(-0.0106, device='cuda:0') tensor(7.5541e-05, device='cuda:0') tensor(0.0082, device='cuda:0')
tensor(-0.0106, device='cuda:0') tensor(7.5541e-05, device='cuda:0') tensor(0.0082, device='cuda:0')
tensor(4.6566e-10, device='cuda:1')
tensor(9.3132e-10, device='cuda:1')
</code></pre>
<p>The issue of rerunning the same line of code and getting another result seems to be really strange.<br>
Was this setup running correctly before or is this a new system?</p>

username: S_Ellis
cooked: <p>Okay, I made a new conda environment and installed the nightly build with <code>conda install pytorch torchvision cudatoolkit=9.2 -c pytorch-nightly</code> as suggested by <a href="https://pytorch.org/get-started/locally/" rel="nofollow noopener">https://pytorch.org/get-started/locally/</a>.</p>
<p>I re-run the script, and I get similar result to before insofar as the outputs do not match:</p>
<pre><code class="lang-auto">tensor(-0.0064, device='cuda:1') tensor(-0.0003, device='cuda:1') tensor(0.0064, device='cuda:1')
tensor(-0.0040, device='cuda:0') tensor(0.0005, device='cuda:0') tensor(0.0056, device='cuda:0')
tensor(-0.0064, device='cuda:0') tensor(-0.0003, device='cuda:0') tensor(0.0064, device='cuda:0')
tensor(3.0040, device='cuda:1')
tensor(3.0049, device='cuda:1')
</code></pre>
<p>And again, if I re-run the last two lines again manually, I get a different output:</p>
<pre><code class="lang-auto">tensor(3.0058, device='cuda:1')
tensor(3.0067, device='cuda:1')
</code></pre>
<p>In this case though, both of these output values change with every run, rather than just the first one.</p>
<p><a class="mention" href="/u/ptrblck">@ptrblck</a> I have never done these kinds of checks before so I don’t know if it was running correctly before. This was the first time I remember that I’ve tried loading a model trained on one GPU onto another.</p>
<p>If both of your set-ups get the correct results, and mine still doesn’t even with the nightly build, what else could it be? I haven’t explicitly set cudnn to deterministic, but I assume that neither of you did either? There must be some mistake I’m making somewhere.</p>