On a cpu device, how to load checkpoint saved on gpu device

Ja-Keoung_Koo: <p>I trained my network on a gpu device and saved checkpoint by <code>torch.save</code></p>
<p>Loading this checkpoint on my cpu device gives an error:</p>
<pre><code class="lang-auto">    raise AssertionError("Torch not compiled with CUDA enabled")
AssertionError: Torch not compiled with CUDA enabled```</code></pre>

apaszke: <p>You can remap the Tensor location at load time using the <code>map_location</code> argument to <code>torch.load</code>. For example this will forcefully remap everything onto CPU:</p>
<pre><code class="lang-python">torch.load('my_file.pt', map_location=lambda storage, location: 'cpu')
</code></pre>
<p>While this will only map storages from GPU0:</p>
<pre><code class="lang-python">torch.load('my_file.pt', map_location={'cuda:0': 'cpu'})
</code></pre>

mromaniuk: <p>I’m trying to load a GPU-trained model onto a CPU with the code you suggested:</p>
<pre><code class="lang-auto">torch.load('my_file.pt', map_location=lambda storage, location: 'cpu')
</code></pre>
<p>… and I get this error:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "net_predict.py", line 146, in &lt;module&gt;
    net = torch.load(f_net, map_location=(lambda storage, location: 'cpu'))
  File "/home/[...]/anaconda2/lib/python2.7/site-packages/torch/serialization.py", line 248, in load
    return _load(f, map_location, pickle_module)
  File "/home/[...]/anaconda2/lib/python2.7/site-packages/torch/serialization.py", line 340, in _load
    tensor = tensor_type._new_with_metadata_file(f, storage)
AttributeError: type object 'str' has no attribute '_new_with_metadata_file'
</code></pre>
<p>(I replaced my username with […])</p>
<p>Any idea what I’m doing wrong?</p>

apaszke: <p>I’m sorry, my bad. This should work:</p>
<pre><code class="lang-python">torch.load('my_file.pt', map_location=lambda storage, loc: storage)
</code></pre>

mromaniuk: <p>It works - brilliant! <img src="//discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=5" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>Out of curiosity: could you explain what this does? I’m not sure how it knows to remap storage to CPU, since the lambda returns the storage it got as an argument.</p>

apaszke: <p>Sure. <code>map_location</code> can be either a dict where the locations corresponding to keys are remaped to their values. Alternatively, we support passing in a function, that will get a CPU storage and its serialized location, and it should return some storage that will replace the CPU one. If you just want to load everything onto the CPU, you can just return the first arugment, but you could do some more crazy stuff like sending all CUDA tensors to the next GPU, by parsing out the original device from the <code>loc</code> argument.</p>

cyyyyc123: <p><a class="mention" href="/u/apaszke">@apaszke</a> Hi! I am sorry to reopen this thread. I have encountered a problem when I used the above method to load a GPU-trained model on CPU mode. The code fragment is:</p>
<pre><code>import torch

encoder = torch.load('encoder.pt', map_location=lambda storage, loc: storage)
decoder = torch.load('decoder.pt', map_location=lambda storage, loc: storage)

encoder.cpu()
decoder.cpu()
</code></pre>
<p>And the error I met was:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/1X/960e33f29babe63783f931cdbec8f0cdf43fd3c4.png" data-download-href="https://discuss.pytorch.org/uploads/default/960e33f29babe63783f931cdbec8f0cdf43fd3c4" title="image.png"><img src="https://discuss.pytorch.org/uploads/default/optimized/1X/960e33f29babe63783f931cdbec8f0cdf43fd3c4_2_690x426.png" width="690" height="426" srcset="https://discuss.pytorch.org/uploads/default/optimized/1X/960e33f29babe63783f931cdbec8f0cdf43fd3c4_2_690x426.png, https://discuss.pytorch.org/uploads/default/original/1X/960e33f29babe63783f931cdbec8f0cdf43fd3c4.png 1.5x, https://discuss.pytorch.org/uploads/default/original/1X/960e33f29babe63783f931cdbec8f0cdf43fd3c4.png 2x" data-dominant-color="0B0A0A"><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">image.png</span><span class="informations">829×513 63.5 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>The full code can be viewed at <a href="https://github.com/Yugnaynehc/pytorch/blob/master/seq2seq-translation/eval.py" rel="noopener nofollow ugc">seq2seq-translation/eval.py</a></p>
<p>How can I load a GPU-trained model on a CPU device (without any GPUs) correctly? Thank you for your great work!</p>

apaszke: <p>Hey, no problem! I only have a couple more questions:</p>
<ol>
<li>What’s your PyTorch version? Do you have <code>torch.__version__</code>? If no, when did you install it?</li>
<li>When did you create that checkpoint?</li>
</ol>

cyyyyc123: <p>Good morning!</p>
<ol>
<li>
<p>The version of my PyTorch is <strong>0.1.9+b46d5e0</strong>. I have compiled the PyTorch from source since I want to try to use half tensor with stateless methods. (You have mentioned it in <a href="https://github.com/pytorch/pytorch/pull/852" rel="nofollow noopener">this pull request</a>. Excellent!)</p>
</li>
<li>
<p>I created the checkpoint about 12 hours before, which also used the  <strong>0.1.9+b46d5e0</strong> version of PyTorch.</p>
</li>
</ol>
<p>Thank you very much!</p>

cyyyyc123: <p>I have uploaded some test data to <a href="https://github.com/Yugnaynehc/pytorch/tree/master/seq2seq-translation" rel="nofollow noopener">my github repo</a>. If you have time maybe you can try it:</p>
<ol>
<li>train a model:<br>
<code>python train_attn.py</code>
</li>
<li>load the model and do some inferences:<br>
<code>python eval.py</code>
</li>
</ol>
<p>May this is useful to provide some information for solving the problem. Thank you!</p>

cyyyyc123: <p><a class="mention" href="/u/apaszke">@apaszke</a> I suggest that it may be the version problem.  I can’t reproduce this error when I use the <strong>0.1.9_2</strong> version of PyTorch. Thanks!</p>

gaoking132: <p>Sorry to reopen the thread.</p>
<p>After running the code:</p>
<blockquote>
<pre><code>params = torch.load(input_file, lambda storage, loc: storage)
</code></pre>
</blockquote>
<p>I met the same problem as Yangyu met before. The error message shows:</p>
<blockquote>
<p>TypeError: set_ received an invalid combination of arguments - got (torch.FloatStorage, int, tuple, tuple), but expected one of:</p>
<ul>
<li>no arguments</li>
<li>(torch.cuda.FloatTensor source)</li>
<li>(torch.cuda.FloatStorage storage)</li>
<li>(torch.cuda.FloatStorage sourceStorage, int storage_offset, int … size)<br>
didn’t match because some of the arguments have invalid types: (!torch.FloatStorage!, int, !tuple!, !tuple!)</li>
<li>(torch.cuda.FloatStorage sourceStorage, int storage_offset, torch.Size size)</li>
<li>(torch.cuda.FloatStorage sourceStorage, int storage_offset, torch.Size size, tuple strides)</li>
</ul>
</blockquote>
<p>I just updated my pytorch to the latest version in the master branch. The version number is 0.1.11+761eef1. Any idea why?</p>
<p>Thanks,<br>
Yaozong</p>

denis_64: <p>Hello, I tried to load a snapshot from gpu-training to run it on CPU-mode, but faced with the same problem, that described above. Of course, tried to use given advice, but there is no effect.</p>
<pre><code class="lang-python">torch.load('./snapshots/cpu_final_snapshot.pth', map_location=lambda storage, loc: storage)
</code></pre>
<p>I have the following traceback:</p>
<pre><code class="lang-python">Traceback (most recent call last):
  File "predict.py", line 39, in &lt;module&gt;
    params = torch.load('./snapshots/cpu_final_snapshot.pth', map_location=lambda storage, loc: storage)
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/torch/serialization.py", line 222, in load
    return _load(f, map_location, pickle_module)
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/torch/serialization.py", line 370, in _load
    result = unpickler.load()
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/torch/cuda/__init__.py", line 279, in __new__
    _lazy_init()
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/torch/cuda/__init__.py", line 96, in _lazy_init
    _check_driver()
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/torch/cuda/__init__.py", line 63, in _check_driver
    raise AssertionError("Torch not compiled with CUDA enabled")
</code></pre>
<pre><code class="lang-python">torch.__version__
'0.1.10_1'
</code></pre>
<p>Would be appreciated any help.</p>

gaoking132: <p>It seems that I found the problem that causes the error of “invalid combination of arguments”.</p>
<p>Yesterday I used the model trained on 0.1.9 version of pytorch, and loaded it to cpu using the latest version of 0.1.11. The error appeared.</p>
<p>Today I retrained the model using the latest version of 0.1.11 and loaded also using the latest version. Everything works.</p>
<p>So I guess that there are inconsistencies between different versions of pytorch models.</p>

Eugenio_Culurciello: <p>We have trained an Alexnet with pytorch examples imagenet (<a href="https://github.com/pytorch/examples/blob/master/imagenet/main.py" rel="nofollow noopener">https://github.com/pytorch/examples/blob/master/imagenet/main.py</a>) and have been struggling to convert the model for use on CPU and for inference only. Here is a solution for AlexNet:<br>
<a href="https://github.com/e-lab/pytorch-toolbox/blob/master/convert-save-load.md" rel="nofollow noopener">https://github.com/e-lab/pytorch-toolbox/blob/master/convert-save-load.md</a><br>
It would be nice to have something more generic…</p>

hazelnutsgz: <p>When I use the torch 1.0.0, the given code will produce the result as the following:<br>
torch.load(‘save/best_BiLSTMCRF_pos_2019-01-10 12-42-50’, map_location=lambda storage, location: ‘cpu’)<br>
Traceback (most recent call last):<br>
File “”, line 1, in <br>
File “/home/jiaxin/.local/lib/python3.6/site-packages/torch/serialization.py”, line 367, in load<br>
return _load(f, map_location, pickle_module)<br>
File “/home/jiaxin/.local/lib/python3.6/site-packages/torch/serialization.py”, line 538, in _load<br>
result = unpickler.load()<br>
File “/home/jiaxin/.local/lib/python3.6/site-packages/torch/_utils.py”, line 135, in _rebuild_tensor_v2<br>
tensor = _rebuild_tensor(storage, storage_offset, size, stride)<br>
File “/home/jiaxin/.local/lib/python3.6/site-packages/torch/_utils.py”, line 129, in _rebuild_tensor<br>
module = importlib.import_module(storage.<strong>module</strong>)<br>
AttributeError: ‘str’ object has no attribute ‘<strong>module</strong>’</p>
<p>Is anything wrong with the new version of PyTorch?</p>

colllin: <p>Had the same thing.  See the comments about using <code>map_location=lambda storage, location: storage</code> instead of <code>'cpu'</code></p>

Zayd: <p>If you want to force the <code>map_location</code> to <code>cpu</code>, you can eliminate the <code>lambda</code> and simply use:</p>
<p><code>torch.load(‘save/best_BiLSTMCRF_pos_2019-01-10 12-42-50’,map_location=‘cpu’)</code></p>
<p>This is discussed in the report for <a href="https://github.com/pytorch/pytorch/issues/9139" rel="nofollow noopener">issue #9139</a>.</p>

Willem: <p>Sorry for reviving this post. I have a closely related question. I want to do the exact same thing, but using the C++ front-end. I.e. I want to save a model, trained using the C++ front-end on GPU, and then load in using the C++ front-end on a CPU device.</p>
<p>It is possible? The documentation on torch::load does not give the map_location? Thanks for any help.</p>

Towsif_Ahamed: <p>torch.load(WEIGHTS_FILE, map_location=torch.device(‘cpu’) )</p>