`torch.load` does not map to GPU as advertised

username: murph213
cooked: <p>I cannot get <code>torch.load</code> and <code>map_location</code> to work as expected.  I have tried three of the suggested methods for loading a model onto the GPU using <code>map_location</code> (from references listed below). The model always ends up on the CPU, despite the documentation seeming to indicate that <code>map_location</code> can load tensors directly onto the GPU.  (It even says this should happen by default on a machine with GPUs)  Here is a minimal working example, followed by further questions.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F

gpu = torch.device("cuda:0")
print(gpu)
</code></pre>
<blockquote>
<p>cuda:0</p>
</blockquote>
<pre><code class="lang-auto"># Build a simple model on the GPU
class Net(nn.Module):
    def __init__(self, indim):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(indim, 5)
        self.fc2 = nn.Linear(5, 1)
        
    def forward(self, t):
        t = F.relu(self.fc1(t))
        t = self.fc2(t)
        return t

    
net = Net(3).to(gpu)

# Force weights to zero so we can later confirm 
# we've loaded the right model
with torch.no_grad():
    net.fc1.weight.fill_(0)
    
print(net.fc1.weight)
</code></pre>
<blockquote>
<p>Parameter containing:<br>
tensor([[0., 0., 0.],<br>
[0., 0., 0.],<br>
[0., 0., 0.],<br>
[0., 0., 0.],<br>
[0., 0., 0.]], device=‘cuda:0’, requires_grad=True)</p>
</blockquote>
<p>Notice my model is on cuda:0</p>
<pre><code class="lang-auto"># Save and delete the model
torch.save(net.state_dict(), 'my_gpu_mod.pth')
del(net)
</code></pre>
<p>Now we try three different ways to load and <code>map_location</code> to the GPU.</p>
<pre><code class="lang-auto"># Load with map_location=device
model = Net(3)
print(model.fc1.weight.max())

model.load_state_dict(torch.load("my_gpu_mod.pth",
                                map_location=gpu))
    
print(model.fc1.weight.device)
print(model.fc1.weight.max())
</code></pre>
<blockquote>
<p>tensor(0.5071, grad_fn=)<br>
cpu<br>
tensor(0., grad_fn=)</p>
</blockquote>
<pre><code class="lang-auto"># Load with map_location=string
model = Net(3)
print(model.fc1.weight.max())

model.load_state_dict(torch.load("my_gpu_mod.pth",
                                map_location="cuda:0"))
    
print(model.fc1.weight.device)
print(model.fc1.weight.max())
</code></pre>
<blockquote>
<p>tensor(0.4118, grad_fn=)<br>
cpu<br>
tensor(0., grad_fn=)</p>
</blockquote>
<pre><code class="lang-auto"># Load with map_location=lambda
model = Net(3)
print(model.fc1.weight.max())

model.load_state_dict(torch.load('my_gpu_mod.pth',
                      map_location=lambda storage, loc: storage.cuda(0)))
    
print(model.fc1.weight.device)
print(model.fc1.weight.max())

</code></pre>
<blockquote>
<p>tensor(0.5070, grad_fn=)<br>
cpu<br>
tensor(0., grad_fn=)</p>
</blockquote>
<p>If <code>map_location</code> does not automatically put things on the GPU, why does it need to be used at all when loading a GPU-trained model on the same machine that trained it?  I can simply call <code>.to</code> afterwards.</p>
<p>References</p>
<ol>
<li><a href="https://pytorch.org/tutorials/beginner/saving_loading_models.html#saving-loading-model-for-inference" class="inline-onebox" rel="noopener nofollow ugc">Saving and Loading Models — PyTorch Tutorials 2.1.1+cu121 documentation</a></li>
<li><a href="https://discuss.pytorch.org/t/on-a-cpu-device-how-to-load-checkpoint-saved-on-gpu-device/349/4" class="inline-onebox">On a cpu device, how to load checkpoint saved on gpu device - #4 by apaszke</a></li>
<li><a href="https://pytorch.org/docs/stable/torch.html#torch.load" class="inline-onebox" rel="noopener nofollow ugc">torch — PyTorch 2.1 documentation</a></li>
</ol>

username: ptrblck
cooked: <p>Could you try to run the code without any available GPUs via <code>export CUDA_VISIBLE_DEVICES=</code> before running the script?<br>
I assume the load method is smart enough to copy the device <code>state_dict</code> back to your CPU model.</p>
<p>EDIT: This issue doesn’t seem to be reproducible using the latest nightly binary, so you might just update. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"></p>

username: murph213
cooked: <p>(1) In fact, I did run without CUDA_VISIBLE_DEVICES.  The machine I’m using only has one GPU and when I run without setting CUDA_VISIBLE_DEVICES, the machine automatically knows to use the only GPU as cuda:0.</p>
<p>(2) I’m not sure I understand this quote</p>
<blockquote>
<p>I assume the load method is smart enough to copy the device <code>state_dict</code> back to your CPU model.</p>
</blockquote>
<p>Actually I am trying to get the <code>map_location</code> argument to load the model from disk, trained on the GPU, directly into the GPU’s VRAM (as advertised in the documentation).  Without having to go through CPU.  So I thought it should not end up on the CPU.</p>
<p>(3) Thank you for the suggestion.  Since my goal is just to understand and I’m wary of nightly binaries, I will pass on them for now.  But to clarify, you are saying, in the newest version, my code above <strong>automatically</strong> loads the model to the GPU’s VRAM without first loading into CPU memory then transferring to the GPU?</p>
<p>Thanks.</p>

username: murph213
cooked: <p>I have made some progress on this issue, by using on <code>torch.cuda.memory_summary</code> rather than <code>mytensor.device</code>.  It turns out that using <code>map_location</code> with options that specify the GPU <em>reserves</em> GPU memory but then stores the model.device as “cpu”.  In particular, it seems that memory is stored on the GPU and then immediately freed.  However, when you use the default <code>map_location</code>, it does not reserve GPU memory at all, GPU memory is never allocated and thus never freed (at least until you call <code>mytensor.to()</code>).</p>
<p>By tracking process CPU memory with <code>psutil</code>, I found that using <code>map_location</code> to GPU causes the total CPU process memory to spike immediately after calling <code>torch.load</code>.  However, if you use the <code>map_location</code> default, the total CPU memory does not increase until after you call <code>mytensor.to()</code>.  In particular, both approaches (with/without <code>map_location</code> and then <code>to()</code>) end up using about the same process CPU memory, but using <code>map_location</code> to GPU gets to that high number after <code>torch.load</code> but <code>map_location</code> to CPU doesn’t get there until after <code>.to</code>.</p>
<p>I’m left wondering whether this was intended behavior.</p>

username: ptrblck
cooked: <p>I still think that you are hunting down two separate issues.<br>
<code>torch.load(..., map_location)</code> will load the tensor or <code>state_dict</code> onto the specified device.<br>
However, since you are piping this operation directly to <code>model.load_state_dict</code>, internally most likely <code>param.copy_</code> will be used, which will then respect the parameters’s device and copy the <code>state_dict</code> parameter to the same device used in the model’s parameter.</p>
<p>Try to separate the calls and check the tensors in the <code>state_dict</code> right after using <code>torch.load</code>.<br>
You should see, that the desired device passed via <code>map_location</code> is used.</p>
<p>If you want to avoid the host and device memory allocation, make sure that the model is on the same device as the <code>state_dict</code> before calling <code>.load_state_dict</code>.</p>

username: murph213
cooked: <p>I understand now; you were right about my source of confusion!</p>
<p>In summary, calling <code>torch.load</code> and <code>.load_state_dict()</code> as a one-liner will put the CPU-trained model on the CPU, even if you specify <code>map_location</code> to a GPU device in <code>torch.load</code>.  However, if one splits these operations into two lines, then <code>torch.load</code> will put the weights in the dictionary onto the desired device.  Then, the second line’s <code>load_state_dict</code> will put the model back onto the CPU, where it was trained from.</p>
<p>The following code and output demonstrate this</p>
<h2>Illustration code</h2>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F

print(f"PyTorch version: {torch.__version__}\n-------")

gpu = torch.device('cuda:0')
PATH = 'todays_cpu_model.pth'

class Net(nn.Module):
    def __init__(self, indim):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(indim, 16)
        self.fc2 = nn.Linear(16, 1)
        
    def forward(self, t):
        t = F.relu(self.fc1(t))
        t = self.fc2(t)
        return t

print("Calling `torch.load` and `model.load_state_dict()` in one line")
model_1 = Net(4)
model_1.load_state_dict(torch.load(PATH, map_location=gpu))
w_1 = next(model_1.parameters()).detach()
print(f"Device of first layer params: {w_1.device}")


print("------\nCalling `torch.load` and `model.load_state_dict()` in TWO DIFFERENT lines")
model_2 = Net(4)
loaded = torch.load(PATH, map_location=gpu)

# access weight in ordered dictionary
w_2a = loaded["fc1.weight"].detach()
print(f"Device of first layer params after `torch.load`: {w_2a.device}")

# now load the state dict and access weight in model object
model_2.load_state_dict(loaded)
w_2b = model_2.fc1.weight.detach()
print(f"Device of first layer params after `.load_state_dict()`: {w_2b.device}")

# Sanity check
# Confirm that the two models are, in fact, equivalent
tests = list()
for kk in model_1.state_dict():
    tests.append(
        torch.equal(model_1.state_dict()[kk].cpu(),model_2.state_dict()[kk].cpu())
    )
    
print(f"----\nTest that the models are identical: {all(tests)}")
</code></pre>
<h2>The output:</h2>
<pre><code class="lang-auto">-------
Calling `torch.load` and `model.load_state_dict()` in one line
Device of first layer params: cpu
------
Calling `torch.load` and `model.load_state_dict()` in TWO DIFFERENT lines
Device of first layer params after `torch.load`: cuda:0
Device of first layer params after `.load_state_dict()`: cpu
----
Test that the models are identical: True
</code></pre>