Best way to save a trained model in PyTorch? [closed]

by

Found this page on their github repo, I’ll just copy-paste the content here.

Recommended approach for saving a model

There are two main approaches for serializing and restoring a model.

The first (recommended) saves and loads only the model parameters:

torch.save(the_model.state_dict(), PATH)

Then later:

the_model = TheModelClass(*args, **kwargs)
the_model.load_state_dict(torch.load(PATH))

The second saves and loads the entire model:

torch.save(the_model, PATH)

Then later:

the_model = torch.load(PATH)

However in this case, the serialized data is bound to the specific classes
and the exact directory structure used, so it can break in various ways when
used in other projects, or after some serious refactors.

Update: see also the Save and Load the Model section from PyTorch tutorial

More Related Contents:

  1. How do I save a trained model in PyTorch?
  2. How to initialize weights in PyTorch?
  3. How do I initialize weights in PyTorch?
  4. Why do we need to call zero_grad() in PyTorch?
  5. How does Pytorch’s “Fold” and “Unfold” work?
  6. What does model.eval() do in pytorch?
  7. PyTorch NotImplementedError in forward
  8. How to understand the term `tensor` in TensorFlow?
  9. Pytorch – RuntimeError: Trying to backward through the graph a second time, but the buffers have already been freed
  10. What’s the difference between torch.stack() and torch.cat() functions?
  11. Is .data still useful in pytorch?
  12. RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same
  13. How do I split a custom dataset into training and test datasets?
  14. RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! when resuming training
  15. PyTorch Binary Classification – same network structure, ‘simpler’ data, but worse performance?
  16. Serializing class instance to JSON
  17. AttributeError: ‘Tensor’ object has no attribute ‘numpy’
  18. Python serialization – Why pickle?
  19. Heroku: slug size too large after installing Pytorch
  20. How do I load a caffe model and convert to a numpy array?
  21. Serialising an Enum member to JSON
  22. How to feed caffe multi label data in HDF5 format?
  23. How to set layer-wise learning rate in Tensorflow?
  24. How do I print the model summary in PyTorch?
  25. How do I check if PyTorch is using the GPU?
  26. Pytorch RuntimeError: The size of tensor a (4) must match the size of tensor b (3) at non-singleton dimension 0
  27. Common use-cases for pickle in Python
  28. looking for source code of from gen_nn_ops in tensorflow
  29. How to fix “ResourceExhaustedError: OOM when allocating tensor”
  30. What difference between pickle and _pickle in python 3?

Leave a Comment