Save and Load model

mjchen611: <p>I have trained a model, I want save it and then reload it and use it to produce the output for new image.</p>
<p>I have found the function  : torch.model(‘path’) ,but when I reload it it always have problem.</p>
<p>Can anyone give me  some suggestions or a simple example?</p>
<p>Thank you so much.</p>

colesbury: <p>Use <code>torch.save</code> and <code>torch.load</code>:</p>
<pre><code class="lang-python">torch.save(model, 'filename.pt')
model = torch.load('filename.pt')
</code></pre>

panovr: <p><a class="mention" href="/u/colesbury">@colesbury</a><br>
I fine-tuning a  resnet-50 model, and load the fine-tuned model as:</p>
<pre><code>model = torch.load('./model_resnet50.pth.tar')

normalize = transforms.Normalize(
    mean=[0.485, 0.456, 0.406],
    std=[0.229, 0.224, 0.225]
)
preprocess = transforms.Compose([
    transforms.Scale(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    normalize
])

img = Image.open(IMG_URL)
img_tensor = preprocess(img)
img_tensor.unsqueeze_(0)
output = model(Variable(img_tensor))
</code></pre>
<p>However, when running the above code, there is an error:<br>
<code>TypeError: 'dict' object is not callable</code></p>

Crazyai: <p>The content in ‘model_resnet50.pth.tar’ is a dict, you should use <code>model = torch.load('./model_resnet50.pth.tar')['state_dict']</code></p>

mjchen611: <p><a class="mention" href="/u/colesbury">@colesbury</a> ok,thank you so much.</p>

rwightman: <p>This covers all the pth files I’ve encountered or created…</p>
<pre><code>        if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
            model.load_state_dict(checkpoint['state_dict'])
        else:
            model.load_state_dict(checkpoint)</code></pre>

_Kanch: <p>Hi ,I’m trying to load model use <strong>model = torch.load(‘<a href="http://cnn4mnist.pt" rel="nofollow noopener">cnn4mnist.pt</a>’)</strong>, but I got error below:</p>
<pre><code>Traceback (most recent call last):
  File "predict.py", line 11, in &lt;module&gt;
    model = torch.load('cnn4mnist.pt')
  File "C:\Users\kanch\Anaconda3\lib\site-packages\torch\serialization.py", line 231, in load
    return _load(f, map_location, pickle_module)
  File "C:\Users\kanch\Anaconda3\lib\site-packages\torch\serialization.py", line 379, in _load
    result = unpickler.load()
AttributeError: Can't get attribute 'CNN4MNIST' on &lt;module '__main__' from 'predict.py'&gt;
</code></pre>
<p>What should I do? I don’t really know where did I wrong?</p>

mjchen611: <p>what is the difference between the file .pkl and the file .pt ?</p>

colesbury: <p>The PyTorch serialization format is built off of pickle (.pkl) but overrides some functionality to handle tensors that share the same backing storage.</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/master/torch/serialization.py" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/master/torch/serialization.py" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/master/torch/serialization.py</a></h4>
<pre><code class="lang-py">import difflib
import inspect
import os
import shutil
import struct
import sys
import torch
import tarfile
import tempfile
import warnings
from contextlib import closing, contextmanager
from ._utils import _import_dotted_name
from ._six import string_classes as _string_classes
if sys.version_info[0] == 2:
    import cPickle as pickle
else:
    import pickle
    import pathlib

DEFAULT_PROTOCOL = 2
</code></pre>

  This file has been truncated. <a href="https://github.com/pytorch/pytorch/blob/master/torch/serialization.py" target="_blank" rel="nofollow noopener">show original</a>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>


rchavezj: <p>I thought the format was in .pth not .pt</p>

saurabhkm: <p>Its arbitrary actually, you can use anything.</p>

woshibali: <p>Is the problem solved? I also encountered the same problem, thank you.</p>

MisBah_Awan: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/1/1237d83331a4367a1d78433973f2e8ab7fbfcd2c.png" data-download-href="https://discuss.pytorch.org/uploads/default/1237d83331a4367a1d78433973f2e8ab7fbfcd2c" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/2X/1/1237d83331a4367a1d78433973f2e8ab7fbfcd2c.png" alt="image" data-base62-sha1="2BaeVJaxi0rRQIfsle4wFrOOmPa" width="666" height="500" data-dominant-color="F6F5F6"><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</span><span class="informations">707×530 20.2 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>

ptrblck: <p>Load the <code>state_dict</code> via:</p>
<pre><code class="lang-python">model = MyModel()
state_dict = torch.load("last_brain1.pth")['state_dict']
model.load_state_dict(state_dict)
</code></pre>

MisBah_Awan: <p>i am stuck here line 154<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/9/933190dda1e4da97fcbd6cbfff6ef5dc9f257dc1.png" data-download-href="https://discuss.pytorch.org/uploads/default/933190dda1e4da97fcbd6cbfff6ef5dc9f257dc1" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/2X/9/933190dda1e4da97fcbd6cbfff6ef5dc9f257dc1_2_690x261.png" alt="image" data-base62-sha1="l08nIAGHjh5Loc7wo9TT9Wbkrw5" width="690" height="261" srcset="https://discuss.pytorch.org/uploads/default/optimized/2X/9/933190dda1e4da97fcbd6cbfff6ef5dc9f257dc1_2_690x261.png, https://discuss.pytorch.org/uploads/default/optimized/2X/9/933190dda1e4da97fcbd6cbfff6ef5dc9f257dc1_2_1035x391.png 1.5x, https://discuss.pytorch.org/uploads/default/original/2X/9/933190dda1e4da97fcbd6cbfff6ef5dc9f257dc1.png 2x" data-dominant-color="F6F5F6"><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</span><span class="informations">1357×515 39.7 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><br>
this solution not work</p>

ptrblck: <p><code>model</code> is still a <code>dict</code> instead of a <code>nn.Module</code> class.<br>
Could you <code>print(model)</code> after running my lines of code?</p>

MisBah_Awan: <p>see<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/b/b797c905188ad3393add8ecc282231c49772ea73.png" data-download-href="https://discuss.pytorch.org/uploads/default/b797c905188ad3393add8ecc282231c49772ea73" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/2X/b/b797c905188ad3393add8ecc282231c49772ea73_2_690x264.png" alt="image" data-base62-sha1="qc8zG3kIuOAGr1cr7GShfTof3B9" width="690" height="264" srcset="https://discuss.pytorch.org/uploads/default/optimized/2X/b/b797c905188ad3393add8ecc282231c49772ea73_2_690x264.png, https://discuss.pytorch.org/uploads/default/optimized/2X/b/b797c905188ad3393add8ecc282231c49772ea73_2_1035x396.png 1.5x, https://discuss.pytorch.org/uploads/default/original/2X/b/b797c905188ad3393add8ecc282231c49772ea73.png 2x" data-dominant-color="F7F6F7"><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</span><span class="informations">1326×509 34.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>

ptrblck: <p><code>model</code> should be initialized as your model class, e.g.:</p>
<pre><code class="lang-python">class MyModel(nn.Module):
    def __init__(self):
        # your layer definitions here

    def forward(self, x):
        # your forward pass here

model = MyModel()
state_dict = torch.load("last_brain1.pth")['state_dict']
model.load_state_dict(state_dict)
</code></pre>

MisBah_Awan: <p>class Net(nn.Module):<br>
    def <strong>init</strong>(self):<br>
        super(Net, self).<strong>init</strong>()<br>
        self.conv1 = nn.Conv2d(3, 32, 3)<br>
        self.pool = nn.MaxPool2d(2, 2)<br>
        self.conv2 = nn.Conv2d(32, 64, 3)<br>
        self.fc1 = nn.Linear(64 <em> 14 </em> 14, 30)<br>
        self.fc2 = nn.Linear(30, 3)<br>
        # self.softmax = nn.Softmax(dim=1)</p>
<p>    def forward(self, x):<br>
        x = self.pool(F.relu(self.conv1(x)))  # output size: [batch_size, 32, 255, 255]<br>
        x = self.pool(F.relu(self.conv2(x)))  # output size: [batch_size, 64, 126, 126]</p>
<p>        x = x.view(-1, 64 <em> 14 </em> 14)  # output size: [batch_size, 64<em>126</em>126]<br>
        x = F.relu(self.fc1(x))<br>
        x = self.fc2(x)<br>
        # x=self.softmax(x)<br>
        return x</p>
<p>x = torch.randn(1, 3, 64, 64)  # (batch size or <span class="hashtag">#of</span> images,channels RGB,width,height)<br>
model = Net()<br>
output = model(x)<br>
model = Net()<br>
state_dict = torch.load(“last_brain1.pth”)[‘state_dict’]<br>
model.load_state_dict(state_dict)<br>
print(model)</p>
<p>here is my code but still not working</p>

ptrblck: <p>What kind error do you get and which line of code raises it?<br>
Could you also show, how you’ve created <code>last_brain1.pth</code>?<br>
Maybe <code>'state_dict'</code> refers to something else than the <code>model.state_dict()</code>?</p>