Note
Click here to download the full example code
Learn the Basics || Quickstart || Tensors || Datasets & DataLoaders || Transforms || Build Model || Autograd || Optimization || Save & Load Model
Save and Load the Model¶
In this section we will look at how to persist model state with saving, loading and running model predictions.
import torch
import torchvision.models as models
Saving and Loading Model Weights¶
PyTorch models store the learned parameters in an internal
state dictionary, called state_dict. These can be persisted via the torch.save
method:
model = models.vgg16(weights='IMAGENET1K_V1')
torch.save(model.state_dict(), 'model_weights.pth')
Downloading: "https://download.pytorch.org/models/vgg16-397923af.pth" to /var/lib/ci-user/.cache/torch/hub/checkpoints/vgg16-397923af.pth
0%| | 0.00/528M [00:00<?, ?B/s]
1%|1 | 6.75M/528M [00:00<00:07, 69.8MB/s]
3%|2 | 15.5M/528M [00:00<00:06, 82.6MB/s]
6%|5 | 31.6M/528M [00:00<00:04, 122MB/s]
9%|9 | 49.0M/528M [00:00<00:03, 145MB/s]
13%|#2 | 66.5M/528M [00:00<00:03, 159MB/s]
16%|#5 | 83.9M/528M [00:00<00:02, 166MB/s]
19%|#9 | 101M/528M [00:00<00:02, 171MB/s]
22%|##2 | 119M/528M [00:00<00:02, 175MB/s]
26%|##5 | 136M/528M [00:00<00:02, 177MB/s]
29%|##9 | 154M/528M [00:01<00:02, 178MB/s]
32%|###2 | 171M/528M [00:01<00:02, 179MB/s]
36%|###5 | 188M/528M [00:01<00:01, 180MB/s]
39%|###8 | 206M/528M [00:01<00:01, 181MB/s]
42%|####2 | 223M/528M [00:01<00:01, 181MB/s]
46%|####5 | 240M/528M [00:01<00:01, 181MB/s]
49%|####8 | 258M/528M [00:01<00:01, 181MB/s]
52%|#####2 | 276M/528M [00:01<00:01, 182MB/s]
56%|#####5 | 293M/528M [00:01<00:01, 182MB/s]
59%|#####8 | 310M/528M [00:01<00:01, 182MB/s]
62%|######2 | 328M/528M [00:02<00:01, 182MB/s]
65%|######5 | 346M/528M [00:02<00:01, 183MB/s]
69%|######8 | 363M/528M [00:02<00:00, 182MB/s]
72%|#######2 | 380M/528M [00:02<00:00, 182MB/s]
75%|#######5 | 398M/528M [00:02<00:00, 182MB/s]
79%|#######8 | 415M/528M [00:02<00:00, 181MB/s]
82%|########1 | 433M/528M [00:02<00:00, 180MB/s]
85%|########5 | 450M/528M [00:02<00:00, 180MB/s]
89%|########8 | 467M/528M [00:02<00:00, 181MB/s]
92%|#########1| 485M/528M [00:02<00:00, 181MB/s]
95%|#########5| 502M/528M [00:03<00:00, 182MB/s]
98%|#########8| 520M/528M [00:03<00:00, 181MB/s]
100%|##########| 528M/528M [00:03<00:00, 175MB/s]
To load model weights, you need to create an instance of the same model first, and then load the parameters
using load_state_dict() method.
model = models.vgg16() # we do not specify ``weights``, i.e. create untrained model
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()
/var/lib/workspace/beginner_source/basics/saveloadrun_tutorial.py:37: FutureWarning:
You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
VGG(
(features): Sequential(
(0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): ReLU(inplace=True)
(2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(3): ReLU(inplace=True)
(4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(6): ReLU(inplace=True)
(7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(8): ReLU(inplace=True)
(9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(11): ReLU(inplace=True)
(12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(13): ReLU(inplace=True)
(14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(15): ReLU(inplace=True)
(16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(18): ReLU(inplace=True)
(19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(20): ReLU(inplace=True)
(21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(22): ReLU(inplace=True)
(23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(25): ReLU(inplace=True)
(26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(27): ReLU(inplace=True)
(28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(29): ReLU(inplace=True)
(30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
)
(avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
(classifier): Sequential(
(0): Linear(in_features=25088, out_features=4096, bias=True)
(1): ReLU(inplace=True)
(2): Dropout(p=0.5, inplace=False)
(3): Linear(in_features=4096, out_features=4096, bias=True)
(4): ReLU(inplace=True)
(5): Dropout(p=0.5, inplace=False)
(6): Linear(in_features=4096, out_features=1000, bias=True)
)
)
Note
be sure to call model.eval() method before inferencing to set the dropout and batch normalization layers to evaluation mode. Failing to do this will yield inconsistent inference results.
Saving and Loading Models with Shapes¶
When loading model weights, we needed to instantiate the model class first, because the class
defines the structure of a network. We might want to save the structure of this class together with
the model, in which case we can pass model (and not model.state_dict()) to the saving function:
torch.save(model, 'model.pth')
We can then load the model like this:
model = torch.load('model.pth')
/var/lib/workspace/beginner_source/basics/saveloadrun_tutorial.py:55: FutureWarning:
You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
Note
This approach uses Python pickle module when serializing the model, thus it relies on the actual class definition to be available when loading the model.
