Saving and loading a model in Pytorch?

Bixqu: <p>If I have a model class and a trainer class. I create an instance of the model and train it.</p>
<pre><code>model = mymodel()
train = trainer.train(model...) 
</code></pre>
<p>How can I save the model to a file, after it has been trained and how can I then reload it and continue training? I searched for this but didn’t get an answer.</p>

Rinku_Jadhav2014: <p><a href="http://pytorch.org/docs/notes/serialization.html" class="onebox" target="_blank" rel="nofollow noopener">http://pytorch.org/docs/notes/serialization.html</a></p>

mratsim: <p><a class="mention" href="/u/rinku_jadhav2014">@Rinku_Jadhav2014</a> unfortunately that tutorial is incomplete to resume training. It will only allow saving a model but it does not save the optimizer, epochs, score, etc.</p>
<p><a class="mention" href="/u/bixqu">@Bixqu</a> You can check the <a href="https://github.com/pytorch/examples/blob/master/imagenet/main.py#L139" rel="nofollow noopener">ImageNet Example line 139</a></p>
<pre><code class="lang-python">        save_checkpoint({
            'epoch': epoch + 1,
            'arch': args.arch,
            'state_dict': model.state_dict(),
            'best_prec1': best_prec1,
            'optimizer' : optimizer.state_dict(),
        }, is_best)
</code></pre>
<p>With</p>
<pre><code class="lang-python">def save_checkpoint(state, is_best, filename='checkpoint.pth.tar'):
    torch.save(state, filename)
    if is_best:
        shutil.copyfile(filename, 'model_best.pth.tar')
</code></pre>
<p>Loading/Resuming from the dictionary is there</p>
<pre><code class="lang-python">    if args.resume:
        if os.path.isfile(args.resume):
            print("=&gt; loading checkpoint '{}'".format(args.resume))
            checkpoint = torch.load(args.resume)
            args.start_epoch = checkpoint['epoch']
            best_prec1 = checkpoint['best_prec1']
            model.load_state_dict(checkpoint['state_dict'])
            optimizer.load_state_dict(checkpoint['optimizer'])
            print("=&gt; loaded checkpoint '{}' (epoch {})"
                  .format(args.resume, checkpoint['epoch']))
        else:
            print("=&gt; no checkpoint found at '{}'".format(args.resume))
</code></pre>

diegslva: <p>Hi! the <strong>best and safe</strong> way to save your model parameters is doing something like this:</p>
<pre><code>  model = MyModel()
  # ... after training, save your model 
  model.save_state_dict('mytraining.pt')

  # .. to load your previously training model:
  model.load_state_dict(torch.load('mytraining.pt'))</code></pre>

mratsim: <p><a class="mention" href="/u/diegslva">@diegslva</a> Unfortunately this has the same issue as the tutorial, it won’t save the epoch and the optimizer state so you can’t resume training which was the OP need.</p>

diegslva: <p><a class="mention" href="/u/mratsim">@mratsim</a>, You’re right! I made a mistake here understanding the question.<br>
I don’t use to do that but, maybe something <code>dirty</code> like that to save you entirely objects:</p>
<pre><code>import copy
import pickle

# model stuff    
model = mymodel()
train = trainer.train(model...)

# copy you entirely object and save it 
saved_trainer = copy.deepcopy(train)
with open(r"my_trainer_object.pkl", "wb") as output_file:
    pickle.dump(saved_trainer, output_file)</code></pre>

ahkarami: <p><a class="mention" href="/u/mratsim">@mratsim</a> &amp; <a class="mention" href="/u/diegslva">@diegslva</a>, when I want to save the trained (i.e., fine tuned) models of ResNet and DenseNet the <code>torch.save(MyModel.state_dict(), './model.pth')</code> method doesn’t work correctly; and when I used the <code>torch.save(MyModel, './model.pth')</code> then the models are saved correctly. It means that when I load my saved models via the first approach, my models don’t give me correct results, however when I use the second approach the results are good. Am I correct? would you please explain why this issue occurred?</p>

smth: <p>when you load the model back again via <code>state_dict</code> method, remember to do <code>MyModel.eval()</code>, otherwise the results will differ.</p>

kellywzhang: <p>Why will the results differ without calling <code>MyModel.eval()</code>?</p>

smth: <p>because your BatchNorm or Dropout layers by default are in <code>train</code> mode on construction.</p>

Prasanna1991: <p>If my model doesn’t use such layers like dropout or batchnorm then it doesn’t make difference to use model or model.eval(), right?</p>

ntomita: <p>You’re right. It matters only when you use those layers, as described in the <a href="http://pytorch.org/docs/master/_modules/torch/nn/modules/module.html#Module.train" rel="nofollow noopener">document</a>. In theory, BN/Dropout should behave differently in evaluation time so you need manually toggle the setting. You could alternatively use model.train(False). Also, make sure to use eval() at validation time.</p>

lan2720: <p>I use <code>.eval()</code> and incorrect either.</p>

ShuokaiPan: <p>HI guys,</p>
<p>I have a question about the behaviour of dropout layer during training and evaluation. I remember reading in a paper that because dropout leave out some units during training. During evaluation, the out going weights of dropout layer need to be reduced an amount corresponding to the dropout rate. For instance, if the dropout rate is 0.5, then the out-going weights need to be reduced by 2, because during evaluation, we effectively have twice the number of units.</p>
<p>So my question is, is this kind of weight scaling mechanism included in the dropout layer in pytorch as well?</p>
<p>Cheers and thanks a lot for your help.<br>
Shuokai</p>

vabh: <p><code>model.eval()</code> takes care of this. However, I think it is scaling the activations and not the weights.</p>

ShuokaiPan: <p>Ok I understand. Thanks for the help.</p>
<p>Cheers</p>

herleeyandi: <p><a class="mention" href="/u/smth">@smth</a> Is it same for <code>model.Train(False)</code> with <code>model.Eval()</code> ?</p>

Hao_Chu: <p>It is true that <code>model.eval()</code> takes care of this. However, it scales when training.</p>
<blockquote>
<p>Furthermore, the outputs are scaled by a factor of 1/(1-p) during training. This means that during evaluation the module simply computes an identity function.</p>
</blockquote>

J_Goldfarb: <p>Newbie question…<br>
Any conventions for filename extensions for saving model and model weights with the following commands?</p>
<p>torch.save(the_model, PATH)</p>
<p>torch.save(the_model.state_dict(), PATH)</p>

smth: <p>we’ve been using <code>.pth</code>, but it’s pretty arbitrary</p>