在本节中，我们将学习如何通过保存、加载和运行模型预测来持久化模型状态。
import torch import torchvision.models as models
保存和加载模型权重
PyTorch 模型将学习到的参数存储在一个名为state_dict的内部状态字典中。这些参数可以通过torch.save方法进行持久化：
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, 398MB/s] 15%|█▌ | 80.1M/528M [00:00<00:01, 422MB/s] 23%|██▎ | 120M/528M [00:00<00:01, 414MB/s] 31%|███ | 164M/528M [00:00<00:00, 429MB/s] 39%|███▊ | 204M/528M [00:00<00:00, 405MB/s] 46%|████▌ | 243M/528M [00:00<00:00, 376MB/s] 53%|█████▎ | 280M/528M [00:00<00:00, 349MB/s] 60%|█████▉ | 314M/528M [00:00<00:00, 353MB/s] 68%|██████▊ | 358M/528M [00:00<00:00, 384MB/s] 76%|███████▌ | 402M/528M [00:01<00:00, 407MB/s] 85%|████████▍ | 446M/528M [00:01<00:00, 423MB/s] 93%|█████████▎| 491M/528M [00:01<00:00, 435MB/s] 100%|██████████| 528M/528M [00:01<00:00, 408MB/s]
要加载模型权重，首先需要创建同一个模型的实例，然后使用load_state_dict()方法加载参数。
在下面的代码中，我们设置weights_only=True以将反序列化期间执行的函数限制为仅加载权重所需的函数。使用weights_only=True是加载权重时的最佳实践。
model = models.vgg16() # we do not specify ``weights``, i.e. create untrained model
model.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) ) )
注意：在进行推理前，请务必调用model.eval()方法，将 dropout 和批量归一化层设置为评估模式。如果不这样做，将会产生不一致的推理结果。
保存和加载带结构的模型
加载模型权重时，我们需要先实例化模型类，因为该类定义了网络的结构。我们可能希望将这个类的结构与模型一起保存，在这种情况下，我们可以将model（而不是model.state_dict()）传递给保存函数：
torch.save(model, 'model.pth')
然后我们可以按照下面的演示加载模型。
保存state_dict被认为是最佳实践。但是，下面我们使用weights_only=False，因为这涉及到加载模型，这是torch.save的旧版用例。
model = torch.load('model.pth', weights_only=False)注意：这种方法在序列化模型时使用 Python pickle 模块，因此在加载模型时，需要确保实际的类定义可用。
