Note
Go to the endto 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#
Created On: Feb 09, 2021 | Last Updated: Sep 25, 2025 | Last Verified: Nov 05, 2024
In this section we will look at how to persist model state with saving, loading and running model predictions.
importtorchimporttorchvision.modelsasmodels
Saving and Loading Model Weights#
PyTorch models store the learned parameters in an internalstate dictionary, calledstate_dict. These can be persisted via thetorch.savemethod:
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] 7%|▋ | 38.0M/528M [00:00<00:01, 397MB/s] 15%|█▍ | 79.0M/528M [00:00<00:01, 416MB/s] 23%|██▎ | 120M/528M [00:00<00:01, 423MB/s] 31%|███ | 161M/528M [00:00<00:00, 425MB/s] 38%|███▊ | 202M/528M [00:00<00:00, 427MB/s] 46%|████▌ | 243M/528M [00:00<00:00, 425MB/s] 54%|█████▎ | 284M/528M [00:00<00:00, 421MB/s] 61%|██████▏ | 324M/528M [00:00<00:00, 424MB/s] 69%|██████▉ | 366M/528M [00:00<00:00, 426MB/s] 77%|███████▋ | 406M/528M [00:01<00:00, 426MB/s] 85%|████████▍ | 447M/528M [00:01<00:00, 425MB/s] 92%|█████████▏| 488M/528M [00:01<00:00, 425MB/s]100%|██████████| 528M/528M [00:01<00:00, 424MB/s]
To load model weights, you need to create an instance of the same model first, and then load the parametersusingload_state_dict() method.
In the code below, we setweights_only=True to limit thefunctions executed during unpickling to only those necessary forloading weights. Usingweights_only=True is considereda best practice when loading weights.
model=models.vgg16()# we do not specify ``weights``, i.e. create untrained modelmodel.load_state_dict(torch.load('model_weights.pth',weights_only=True))model.eval()
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 callmodel.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 classdefines the structure of a network. We might want to save the structure of this class together withthe model, in which case we can passmodel (and notmodel.state_dict()) to the saving function:
torch.save(model,'model.pth')
We can then load the model as demonstrated below.
As described inSaving and loading torch.nn.Modules,savingstate_dict is considered the best practice. However,below we useweights_only=False because this involves loading themodel, which is a legacy use case fortorch.save.
model=torch.load('model.pth',weights_only=False)
Note
This approach uses Pythonpickle module when serializing the model, thus it relies on the actual class definition to be available when loading the model.
Related Tutorials#
Total running time of the script: (0 minutes 5.528 seconds)
