Saving and Loading Trained Model Parameters

username: ahmed
cooked: <p>I want to train from scratch the AlexNet model:</p>
<pre><code class="lang-auto">model = models.alexnet(pretrained=False)
num_features = model.classifier[6].in_features
features = list(model.classifier.children())[:-1] # Remove last layer
features.extend([nn.Linear(num_features, 4)]) # Add our layer with 4 outputs
model.classifier = nn.Sequential(*features) # Replace the model classifier

for param in model.features.parameters():
    param.require_grad = True
</code></pre>
<p>I want to save the parameters/hyperparmeters (weights, bias, model values/structure etc.).</p>
<p>This is the training function that I am using:</p>
<pre><code class="lang-auto">def train(model, dataloaders, criterion1, criterion2, optimizer, num_epochs):
    since = time.time()
    train_loss = []
    val_acc_history = []
    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0
    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch + 1, num_epochs))
        print('-' * 10)
        for phase in ['train', 'validation']:
            if phase == 'train':
                model.train() 
            else:
                model.eval()  
            running_loss = 0.0
            running_corrects = 0 
            for images, target in dataloaders[phase]:
                images = images.to(device)
                labels = target['label'].to(device)
                norm = target['norm_box'].to(device)
                
                optimizer.zero_grad()
                with torch.set_grad_enabled(phase == 'train'):
                    out_bbox = model(images)
                    loss_bbox = criterion2(out_bbox, norm)
                    loss = loss_bbox

                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                running_loss += loss.item() * images.size(0)
                running_corrects += IoU(norm, out_bbox, 0.3, batch_size)

            epoch_loss = running_loss / len(dataloaders[phase].dataset)
            if phase == 'train':    
                train_loss.append(epoch_loss)
            epoch_acc = running_corrects / len(dataloaders[phase].dataset)

            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
            
            if phase == 'validation' and epoch_acc &gt; best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())
            if phase == 'validation':
                val_acc_history.append(epoch_acc)

        print()
    time_elapsed = time.time() - since
    print('Training complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
    print('Best val Acc: {:4f}'.format(best_acc))
    torch.save(best_model_wts, 'weights/best_wts.pth')
    return model, train_loss, val_acc_history
</code></pre>
<p>When I load the <code>state_dict</code>:</p>
<pre><code class="lang-auto">wts = model.load_state_dict(torch.load('weights/best_wts.pth))
</code></pre>
<p>I get the message:</p>
<pre><code class="lang-auto">IncompatibleKeys(missing_keys=[], unexpected_keys=[])
</code></pre>
<p>And when trying to look through the state_dict:</p>
<pre><code class="lang-auto">print(list(wts.keys()))
</code></pre>
<p>I get the error:</p>
<pre><code class="lang-auto">AttributeError: 'IncompatibleKeys' object has no attribute 'keys'
</code></pre>
<p>Any advice would be greatly appreciated.</p>

username: ptrblck
cooked: <p>This message just states, that <strong>no</strong> incompatible keys were found, so you are good to go.<br>
As this message might be confusing, we have a discussion <a href="https://github.com/pytorch/pytorch/issues/20128" rel="nofollow noopener">here</a> to disable it if no incompatible keys were found.</p>

username: ahmed
cooked: <p>Thanks for the reply.</p>
<p>Could you please also confirm if this method does indeed save the parameters/hyperparameters (weights, bias, model values/structure etc.) of the model?<br>
I hope that makes sense.</p>

username: ptrblck
cooked: <p>Your current code will only save the <code>model.state_dict</code>, i.e. all parameters of your model.<br>
The model structure will not be saved, so you should always keep the model definition close to the checkpoint.<br>
Also, if you would like to use this checkpoint for finetuning, I would recommend to store the <code>optimizer.state_dict()</code>, since some optimizers have internal parameters (e.g. Adam), which would be otherwise reset to their default values.<br>
Have a look at the <a href="https://github.com/pytorch/examples/blob/1de2ff9338bacaaffa123d03ce53d7522d5dcc2e/imagenet/main.py#L249" rel="nofollow noopener">ImageNet example</a> to see, how the checkpoint is created.</p>

username: ahmed
cooked: <p>Thank you very much for the explanation and the link.</p>

username: Brando_Miranda
cooked: <p>is there a reason why the standard way doesn’t work for you:</p>
<aside class="onebox stackexchange">
  <header class="source">
      <a href="https://stackoverflow.com/questions/42703500/best-way-to-save-a-trained-model-in-pytorch" target="_blank" rel="nofollow noopener">stackoverflow.com</a>
  </header>
  <article class="onebox-body">
      <a href="https://stackoverflow.com/users/5352399/wasi-ahmad" target="_blank" rel="nofollow noopener">
    <img alt="Wasi Ahmad" src="https://lh3.googleusercontent.com/-Rk_dLwN-xPY/AAAAAAAAAAI/AAAAAAAAAC8/DtCya1W9lX4/photo.jpg?sz=128" class="thumbnail onebox-avatar" width="60" height="60">
  </a>
<h4>
  <a href="https://stackoverflow.com/questions/42703500/best-way-to-save-a-trained-model-in-pytorch" target="_blank" rel="nofollow noopener">Best way to save a trained model in PyTorch?</a>
</h4>

<div class="tags">
  <strong>python, serialization, deep-learning, pytorch</strong>
</div>

<div class="date">
  asked by
  
  <a href="https://stackoverflow.com/users/5352399/wasi-ahmad" target="_blank" rel="nofollow noopener">
    Wasi Ahmad
  </a>
  on <a href="https://stackoverflow.com/questions/42703500/best-way-to-save-a-trained-model-in-pytorch" target="_blank" rel="nofollow noopener">07:06PM - 09 Mar 17 UTC</a>
</div>

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

<hr>
<p>see answer above:</p>
<p>ref: <a href="https://github.com/pytorch/pytorch/blob/761d6799beb3afa03657a71776412a2171ee7533/docs/source/notes/serialization.rst" rel="nofollow noopener">https://github.com/pytorch/pytorch/blob/761d6799beb3afa03657a71776412a2171ee7533/docs/source/notes/serialization.rst</a></p>
<hr>
<h1>Recommended approach for saving a model</h1>
<p>There are two main approaches for serializing and restoring a model.</p>
<p>The first (recommended) saves and loads only the model parameters:</p>
<pre><code>torch.save(the_model.state_dict(), PATH)
</code></pre>
<p>Then later:</p>
<pre><code>the_model = TheModelClass(*args, **kwargs)
the_model.load_state_dict(torch.load(PATH))
</code></pre>
<p>The second saves and loads the entire model:</p>
<pre><code>torch.save(the_model, PATH)
</code></pre>
<p>Then later:</p>
<pre><code>the_model = torch.load(PATH)
</code></pre>
<p>However in this case, the serialized data is bound to the specific classes<br>
and the exact directory structure used, so it can break in various ways when<br>
used in other projects, or after some serious refactors.</p>