Loading weights for CPU model while trained on GPU

username: jtremblay
cooked: <p>This is not a very complicated issue, but I am not sure what is the best way to load the weights into the cpu when the model was trained on a GPU, thus here is my solution:</p>
<pre><code class="lang-python">model = torch.load('mymodel')
self.model = model.cpu().double()
</code></pre>
<p>I am not sure if this should be a bug, also this discussion is related <a href="https://discuss.pytorch.org/t/on-a-cpu-device-how-to-load-checkpoint-saved-on-gpu-device/349/8">link</a>.</p>

username: fmassa
cooked: <p>There is <a href="http://pytorch.org/docs/torch.html?highlight=load#torch.load">an option in <code>torch.load</code></a> to map the parameters from one specified device to another.<br>
<a href="https://discuss.pytorch.org/t/on-a-cpu-device-how-to-load-checkpoint-saved-on-gpu-device/349/4?u=fmassa">As pointed out in the link you send</a>, there is a way of forcing all GPU tensors to be in CPU while loading, which I copy here:</p>
<pre><code class="lang-python">torch.load('my_file.pt', map_location=lambda storage, loc: storage)
</code></pre>

username: achaiah
cooked: <p>Quick question… this only seems to work when the model is trained on one GPU. If I train my model on multiple GPUs, save it, and then try to load on the CPU I get this error: <code>KeyError: 'unexpected key "module.conv1.weight" in state_dict'</code>  Is there something different that needs to happen during saving/loading for multiple GPUs?</p>

username: fmassa
cooked: <p>This is a different issue, and is related to <a href="https://discuss.pytorch.org/t/solved-keyerror-unexpected-key-module-encoder-embedding-weight-in-state-dict/1686">[solved] KeyError: 'unexpected key "module.encoder.embedding.weight" in state_dict'</a></p>

username: achaiah
cooked: <p>Perfect, thanks for the clarification.  I ended up using the conversion suggested in your linked post.</p>

username: jadore801120
cooked: <aside class="quote no-group quote-modified" data-username="fmassa" data-post="2" data-topic="1032">
<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/fmassa/48/58_2.png" class="avatar"> fmassa:</div>
<blockquote>
<p>torch.load(‘my_file.pt’, map_location=lambda storage, loc: storage)</p>
</blockquote>
</aside>
<p>Thanks for the suggestion!</p>

username: tralgu
cooked: <p>Wait, I don’t get it…<br>
So in the lambda:</p>
<p><code>torch.load('my_file.pt', map_location=lambda storage, loc: storage)</code></p>
<p><code>storage</code> and <code>loc</code> are supposed to be replaced by the variable I want to store the value in and the target location (CPU or GPU number) respectively, right? But how do I specify the location? Are there any ‘keywords’ to do so?</p>

username: Peter_Ham
cooked: <p>I have a related question, I have a shared model trained on GPU, and another process needs this model for inference on CPU. So I use a shared model and use the following command to load this shared model</p>
<pre><code>    cpu_model.load_state_dict(gpu_model.cpu().state_dict())
</code></pre>
<p>however, this won’t work and returns CUDA error(3), initialization error, what happened?</p>

username: Marat
cooked: <p>If I load model with your hack and set model.train(‘True’) and trying even inference it fails. So it do not work for all cases (finetune on CPU after training on GPU not working).</p>

username: TSLsun
cooked: <p>Not sure about your situation.<br>
However, maybe try the following code, which works for me:</p>
<pre><code class="lang-auto">use_cuda = torch.cuda.is_available()
DEVICE = torch.device('cuda' if use_cuda else 'cpu')   # 'cpu' in this case

cpu_model = your_model()
cpu_model.load_state_dict(torch.load(path_to_your_saved_gpu_model, map_location=DEVICE)
</code></pre>

username: A_Rza_SH
cooked: <p>thanks, worked for me</p>

username: prashanth
cooked: <p>sir in this post I didn’t understand what to place in the place of your_model() can you please show it using example .<br>
thank you waiting for the response…</p>