30分钟快速入门Pytorch

0、LEARN THE BASICS

作者：

大多数的机器学习流程包括数据处理、创造模型、优化模型参数、保存训练模型。这个教程教你通过Pytorch实现机器学习流程，你可以在本教程中了解更多机器学习的概念。

我们将使用FashionMNIST 数据集去训练一个神经网络预测输入图像是否属于接下来的任意一类：T恤/上衣、裤子、套头衫、连衣裙、外套、凉鞋、衬衫、运动鞋、包或踝靴。

这个教程假定你熟悉基本的Python和深度学习的概念。

0.0 运行教程代码

你可以运行这个代码以下面的方式

1）在云端：这是最简单的方式开始！每个部分在顶端有一个“运行在微软学习”的链接，能够在一个完全托管的环境中打开一个集成的笔记本，其中包含代码。

2）本地：这个选择要求你先去建立Pytorch和Torchvision在你的本地机器上。下载这个笔记或者复制代码到你最喜欢的IDE下。

0.1 怎样使用向导？

如果你熟悉其他的深度学习框架，首先检查这个0.Quickstart 去使你自己快速熟悉Pytorch的API。

如果你刚刚接触深度学习的框架，直奔手册的第一部分，按部就班的学习！

0.Quickstart

1.Tensors

2.Datasets and DataLoaders

3.Transforms

4.Build Model

5.Automatic Differentiation

6.Optimazation Loop

7.Save,Load and Use Model

1. TENSORS

Tensors是一个特别的数据结构，与数组和矩阵相似。在Pytorch中,我们使用tensor去编码模型的输入和输出，还有模型的参数。

Tensors与Numpy的数组类似，除了Tensors能够在GPU或者其他硬件加速器上运行。事实上，tensors和numpy的数组共享同一份底层内存，无需复制数据。Tensors对自动求导进行了优化，我们将会在自动求导部分看到更多细节。如果你熟悉数组，你将可以熟练使用Tensor，如果不，请跟随！

1.1 初始化数组

Tensors能够被以各种各样的方式初始化。请看接下来的例子！

import torch
import numpy as np

1）直接从数据

Tensors能够被直接创建从数据。数据类型被自动推到。

data = [[1,2],[3,4]]
x_data = torch.tensor(data)
print(x_data)

2）从Numpy数组

Tensors能够被创建根据numpy数组

np_array = np.array(data)
x_np = torch.from_numpy(np_array)
print(x_np)

3）从其他Tensor

除非显示重写，否则新的张量将会保留参数张量的属性（形状、数据类型）。

x_rand = torch.rand_like(x_data, dtype=torch.float) # overrides the datatype of x_data

4）随机或者常量值

shape是一个tensor维数的元组。在接下来的函数中，它决定输出张量的维数。

shape = (2,3,)
rand_tensor = torch.rand(shape)
ones_tensor = torch.ones(shape)
zeros_tensor = torch.zeros(shape)
​
print(f"Random Tensor: \n {rand_tensor} \n")
print(f"Ones Tensor: \n {ones_tensor} \n")
print(f"Zeros Tensor: \n {zeros_tensor}")

1.2 张量的属性

张量属性描述它们的形状、数据类型、和他们存储的设备。

tensor = torch.rand(3,4)
​
print(f"Shape of tensor: {tensor.shape}")
print(f"Datatype of tensor: {tensor.dtype}")
print(f"Device tensor is stored on: {tensor.device}")

1.3 在Tensor上的操作

超过一百个tensor操作，包括算数、线性代数、矩阵操作、采样等被详细介绍

这些操作都能被运行在GPU（比CPU的速度快很多），如果你使用Colab,分配一个GPU通过点击Runtime-Change runtime type-GPU.

默认来说，Tensors被创建在CPU上。在检查GPU是可用的之后，我们需要去明确的使用.to方法移动Tensors到GPU上，请记住，跨设备复制大张量可能会耗费大量时间和内存！

# We move our tensor to the GPU if available
if torch.cuda.is_available():
    tensor = tensor.to("cuda")

尝试列表中的一些操作。如果你熟悉NumPy API，你会发现Tensor API很容易使用。

1.4 标准的Numpy索引和切片

tensor = torch.ones(4, 4)
print(f"First row: {tensor[0]}")
print(f"First column: {tensor[:, 0]}")
print(f"Last column: {tensor[..., -1]}")
tensor[:,1] = 0
print(tensor)

1.5 聚合tensors

你能使用torch.cat沿着给定维度去连接一系列张量，torch.stack，另外一个连接张量的命令，不同于Torch.cat。

t1 = torch.cat([tensor, tensor, tensor], dim=1)
print(t1)

1.6 算数运算

1.7 就地操作

存储结果到操作数中的操作被叫做就地操作。被_后缀表示。

print(f"{tensor} \n")
tensor.add_(5)
print(tensor)

2、DATASETS & DATALOADERS

预处理数据样本的代码混乱并且难以维持。理想情况下，我们希望数据集代码与模型训练代码分离，增强可读性与模块性。

Pytorch提供两种数据原语：torch.utils.data.DataLoader 和torch.utils.data.Dataset，允许你使用预加载的数据集和你自己的数据。

Dataset存储数据样本和对应的标签，DataLoader打包一个迭代器在Dataset周围，使访问数据更容易。

Pytorch的域库提供了大量的预加载的数据集，例如FashionMNIST,是torch.utils.data.Dataset的子集，实现特定于特定数据的函数。他们可以用于原形和基准测试你的模型。你能发现他们在这：

图像数据：Datasets — Torchvision 0.12 documentation (pytorch.org)

为了建立数据集，Torchvision提供一些torchvision.datasets内置的数据库模型，还有建立自己数据集的实用程序类。

内置数据类型：

全部的数据集是torch.utils.data.Dataset的子类，他们有getitem和len方法。因此，他们全部都能被传递到torch.utils.data.DataLoader中，能够使用torch.multiprocessing并行加载多个样本。例如：

全部的数据集有相似的API。他们全部都有两个共同的参数，transform 和target_transform 去变换输入和输出。你也能够创造你自己的数据集使用下面这些基类。

2.1 加载一个数据集

Torchvision是一个视觉工具包，这个包独立于Pytorch,需要通过pip install torchvision安装

torchvision主要包含三部分：

models：提供深度学习中各种经典网络的网络结构以及训练好的模型，包括AlexNet,VGG,ResNet,Inception等

datasets:提供常用的数据集加载，设计上都是继承torch.utils.data.Dataset

transforms:提供常用的数据预处理操作

这是一个例子-怎样从TorchVision中加载Fashion-MNIST数据集。Fashion-MNIST是一个Zalando's文章的数据集，由60000个训练集例子和10000个测试集例子组成。每个例子由28x28的灰度图像和对应标签组成。

我们加载数据集用接下来的参数：

root是训练集或者测试集存储的位置

train指的是训练或者测试数据集

download=True如果它不可获得从root,加载数据从互联网上。

transform 和 target_transform指的是特征和标签转变。

import torch
from torch.utils.data import Dataset
from torchvision import datasets
from torchvision.transforms import ToTensor
import matplotlib.pyplot as plt
​
​
training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor()
)
​
test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor()
)

2.2 迭代和可视化数据集

我们像list一样手动的使用索引Datasets：training_data[index]，我们使用matplotlib去可视化我们训练集中的一些样本。

labels_map = {
    0: "T-Shirt",
    1: "Trouser",
    2: "Pullover",
    3: "Dress",
    4: "Coat",
    5: "Sandal",
    6: "Shirt",
    7: "Sneaker",
    8: "Bag",
    9: "Ankle Boot",
}
figure = plt.figure(figsize=(8, 8))
cols, rows = 3, 3
for i in range(1, cols * rows + 1):
    sample_idx = torch.randint(len(training_data), size=(1,)).item()
    img, label = training_data[sample_idx]
    figure.add_subplot(rows, cols, i)
    plt.title(labels_map[label])
    plt.axis("off")
    plt.imshow(img.squeeze(), cmap="gray")
plt.show()

2.3 为自己的文件创建一个自定义数据集

一个自定义数据集类一定要实现三个功能init,len,getitem，看一下这些实现。FashionMNIST图像被存储在目录img_dir下面，他们的标签被存储在一个CSV文件annotations_file中。

在下面的部分，我们分节讲解每部分函数的功能。

import os
import pandas as pd
from torchvision.io import read_image

class CustomImageDataset(Dataset):
    def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):
        self.img_labels = pd.read_csv(annotations_file)
        self.img_dir = img_dir
        self.transform = transform
        self.target_transform = target_transform

    def __len__(self):
        return len(self.img_labels)

    def __getitem__(self, idx):
        img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
        image = read_image(img_path)
        label = self.img_labels.iloc[idx, 1]
        if self.transform:
            image = self.transform(image)
        if self.target_transform:
            label = self.target_transform(label)
        return image, label

2.3.1 init

init函数在实例化Dataset对象时运行一次。我们初始化包含图像目录、注释、还有两种转换。

标签labels.csv文件格式如下：

tshirt1.jpg, 0
tshirt2.jpg, 0
......
ankleboot999.jpg, 9

def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):
    self.img_labels = pd.read_csv(annotations_file)
    self.img_dir = img_dir
    self.transform = transform
    self.target_transform = target_transform

2.3.2 len

len函数返回我们数据集样本的数量！

def __len__(self):
    return len(self.img_labels)

2.3.3 getitem

函数加载并且返回数据集中索引为idx的样本。依赖于索引，它指出图像在磁盘中的位置，使用read_image将其转换为tensor类型，取对应的标签从self.img_labels数据中，应用transform函数并且返回tensor图像和对应的标签元组。

def __getitem__(self, idx):
    img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
    image = read_image(img_path)
    label = self.img_labels.iloc[idx, 1]
    if self.transform:
        image = self.transform(image)
    if self.target_transform:
        label = self.target_transform(label)
    return image, label

2.4 准备数据用DataLoaders加载训练

数据集一次检索我们数据集的特征和标签。当训练模型时，我们通常以小批量的方式传递样本，每个历元去重新打乱数据减少模型的过拟合，使用Python的多处理加速数据检索。

DataLoader是一个迭代器，以简单的API为我们抽象了复杂性。

from torch.utils.data import DataLoader

train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=64, shuffle=True)

2.5 通过DataLoader进行迭代

我们已经加载数据集到DataLoader中，并且可以根据需要遍历数据集。每个迭代返回一批训练特征和训练标签。因为我们指明shuffle = True,在我们迭代全部的批次后，数据被洗牌。

# Display image and label.
train_features, train_labels = next(iter(train_dataloader))
print(f"Feature batch shape: {train_features.size()}")
print(f"Labels batch shape: {train_labels.size()}")
img = train_features[0].squeeze()
label = train_labels[0]
plt.imshow(img, cmap="gray")
plt.show()
print(f"Label: {label}")

2.6 准备数据以便使用数据加载程序进行培训

这个Dataset一次取出一个特征和标签。当训练一个模型的时候，我们想要以“minibatches”的形式传递样本，每次epoch弄乱数据去减少模型过拟合，并且使用Python的multiprocessing去加速数据取出。

Dataloader是一个简单的API。

遍历数据加载器

我们已经加载数据到Dataloader中，并且按照需求迭代数据集。每个迭代器返回一个批次的train_features和train_labels。因为我们标记shuffle=True,在我们迭代全部的batches之后，数据被打乱，要对数据加载顺序进行更细粒度的控制，请查看采样器。

3、转换

数据并不总是以训练机器学习算法所需的最终处理形式出现。我们使用转换对数据进行一些操作，使其适合于训练。

全部的TorchVision数据集有两个参数，transform去修改特征并且target_transform去修改标签，接受包含转换逻辑的可调用项。torchvision.transforms模型提供了几种常用的现成的转换。

FashionMNIST特征是PIL图像格式，标签是整数格式。为了训练，我们需要标准的Tensor类型的特征，标签是独热编码Tensor。为了实现这些转换，我们使用ToTensor和Lambda。

import torch
from torchvision import datasets
from torchvision.transforms import ToTensor, Lambda

ds = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor(),
    target_transform=Lambda(lambda y: torch.zeros(10, dtype=torch.float).scatter_(0, torch.tensor(y), value=1))
)

3.1 ToTensor()

ToTensor()将一个PIL图像或者Numpy数组转换为FloatTensor。并且将像素值的范围缩放到0-1之间。

3.2 Lambda Transforms

Lambda transforms 应用任何使用者定义的Lambda函数。在这，我们定义函数将整数转换为独热编码的Tensor。它首先创造一个大小为10（数据集中标签的类别）的0Tensor，然后应用scatter函数分配一个value=1在被给标签y的索引。

scatter_()

4、搭建神经网络

神经网络由操作数据的模型或者层组成。torch.nn提供需要搭建神经网络的全部模块。pytorch中的每个模块都是nn.Module的子类。一个神经网络是一个由其他模型组成的模块，这种嵌套结构允许轻松构建和管理复杂的体系结构。

在接下来的部分，我们将搭建一个神经网络去分类FashionMNIST数据集。

import os
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

4.1 得到训练设备

如果条件允许，我们想要训练我们的模型在一个类似GPU的硬件加速器上。我们检查一下torch.cuda是否是可以获得的，否则我们继续使用CPU。

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using {device} device")

4.2 定义一个类

我们通过继承nn.module定义我们自己的神经网络,并且初始化这个网络层在init。每个nn.Module的子类在forward方法中实现输入数据的处理。

class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10),
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

4.3 将搭建好的网络模型传送到device中

我们创造一个NeuralNetwork的例子，并且将它移动到device中，并且输出它的结构。

model = NeuralNetwork().to(device)
print(model)

为了使用这个模型，我们传入输入数据。这将执行模型的正向操作，以及一些后台操作（不要直接调用model.forward()）。对输入调用模型，返回一个十维的张量，包含每个类的预测值。我们通过nn.Softmax得到预测概率。

X = torch.rand(1, 28, 28, device=device)
logits = model(X)
pred_probab = nn.Softmax(dim=1)(logits)
y_pred = pred_probab.argmax(1)
print(f"Predicted class: {y_pred}")

4.4 模型层

让我们将FashionMNIST模型的层拆开。为了去解释这个，我们将会

举一个三张图象的例子并且看一看当我们将它传进网络会发生什么。

4.4.1 nn.Flatten

我们初始化nn.Flatten层去将图片转化为一个连续数组。

4.4.2 nn.Linear

线性层是一个模块，它使用存储的权重和偏差对输入应用线性变换。

4.4.3 nn.ReLu

非线性激活就是创造复杂的映射在输入和输出数据之间。在应用线性变换之后，他们被引用去引进非线性，帮助神经网络学习各种各样的现象。

在这，我们使用nn.ReLu在我们的线性层之间，但是这也有其他的非线性函数。

4.4.4 nn.Sequential

nn.Sequential 是一个模型的顺序容器。数据按顺序被模型传递。你可以使用顺序容器组合一个快速的网络seq_modules。

4.4.5 nn.Softmax

4.5 模型参数

神经网络内部的一些层被参数化，在训练过程中有对应的权重和偏置需要被优化。nn.Module的子类自动的追踪对象定义内的所有字段，摈弃给所有参数可以访问通过使用parameters()或者named_parameters()方法。

在本例中，我们迭代每个参数，并打印其大小及其值的预览。

5、自动求导

当训练神经网络的时候，最经常使用的算法是回传。在这部分中，参数被调整根据损失函数对给定参数的梯度。

为了计算这些梯度，Pytorch建立了一个微分引擎叫做torch.autograd。它支持任何计算图形的梯度自动求导。

考虑到最简单的一层神经网络，输入x，参数w，偏置b，和一些损失函数。能够在Pytorch中定义为如下格式。

import torch

x = torch.ones(5)  # input tensor
y = torch.zeros(3)  # expected output
w = torch.randn(5, 3, requires_grad=True)
b = torch.randn(3, requires_grad=True)
z = torch.matmul(x, w)+b
loss = torch.nn.functional.binary_cross_entropy_with_logits(z, y)

5.1 张量、函数、计算图

代码被定义为如下的计算图

 

在这个网络中，w和b是参数，需要我们去优化。因此，我们需要去计算损失函数对这些变量的梯度。为了做这个，我们设置这些tensors的requires_grad属性。

注意：你能够设置requires_grad当创建一个Tensor的时候，或者稍后使用x.requires_grad_(True)方法。

一个我们应用到tensors上去构造计算图的函数实际上是一个Function类。这个类知道怎么样去计算函数在前向中，并且也知道怎么样去计算微分在回传中。对反向传播函数的引用存储在张量的grad_fn属性中。您可以在文档中找到有关函数的更多信息。

print(f"Gradient function for z = {z.grad_fn}")
print(f"Gradient function for loss = {loss.grad_fn}")

5.2 计算梯度

为了去优化神经网络中的权重参数，我们需要去计算损失函数对这些参数的微分。为了计算这些微分，我们需要调用loss.backward()，并且获取值从w.grad和b.grad。

loss.backward()
print(w.grad)
print(b.grad)

注意：我们仅仅获得计算图中叶节点的梯度属性，我们已经设置requires_grad属性为True，对于图中的其他点，梯度不是可获得的。

我们在一个图像中仅仅执行一次backward()，如果我们需要做几次backward,我们需要设置retain_graph = True。

5.3 禁用梯度跟踪

默认来讲，全部的张量都带有初始属性requires_grad = True去追踪他们的计算历史并且支持梯度计算。然而，有的时候我们不需要这样做，例如，当我们训练好模型之后，我们仅仅想要将输入数据传递进去，我们仅仅需要前向计算。这种情况，我们能够通过将我们的代码放置在torch.no_grad()中去停止追踪计算。

z = torch.matmul(x, w)+b
print(z.requires_grad)

with torch.no_grad():
    z = torch.matmul(x, w)+b
print(z.requires_grad)

另一个方法是去使用detach()去实现：

z = torch.matmul(x, w)+b
z_det = z.detach()
print(z_det.requires_grad)

这有很多你想要让梯度追踪失效的原因：

5.4 计算图的更多信息

概念上讲，自动求导记录了全部的数据和全部的操作。

• 在你的神经网络中标记一些参数作为冻结参数。这是微调预训练网络的常见场景。

• 当你仅仅需要做前向传输时，加速计算，因为不追踪梯度的计算更加高效。

6、优化模型参数

既然我们已经有了模型、数据，是时候去训练、验证并且测试我们的模型通过在我们的数据上优化参数。训练一个模型是一个迭代的过程，在每个迭代中，模型猜测输出，计算猜测的误差。收集误差对于参数的导数，使用梯度下降法优化参数。

6.1 参数调整

超参数调整可以在平均模型和高精度模型之间产生差异。通常情况像选择简单的学习率或者改变网络层大小这种操作可能会对我们的模型产生巨大的影响。

幸运的是，有工具能够帮助我们发现最好的参数组合。Ray Tune是一个标准的分布式调参工具。Ray Tune包括最新的超参数搜索算法，与TensorBoard和其他分析库集成，并通过Ray的分布式机器学习引擎支持分布式训练。

在这个教程中，我们将展示如何集成Ray Tune到你的Pytorch训练流程中。我们将扩展Pytorch教程去训练一个CIFAR10图像分类器。

正如你看到的，我们仅仅需要去添加一些简单的修改。我们尤其需要：

1.打包数据加载和训练的函数

2.修改网络参数

3.添加检查点

4.定义模型调参的搜索空间

为了运行这个教程，请确信接下来的包已经被安装：

ray[tune]:分布式的超参数调整库

torchvision:数据格式转换

Setup /Imports

6.1 超参数

超参数是可调整的参数，可以用于控制模型优化过程。不同的超参数值会影响模型训练和收敛速度。

我们定义如下的超参数：

• Epochs的大小：整个数据集的迭代次数

• Batch size:更新参数之前通过网络传播的数据样本数

• Learning Rate: 在每个批次/时期更新模型参数的数量。较小的值会导致学习速度较慢，而较大的值可能会导致训练期间出现不可预测的行为。

learning_rate = 1e-3
batch_size = 64
epochs = 5

6.2 优化循环

一旦我们设定了超参数，我们就可以通过一个优化循环来训练和优化我们的模型。优化循环的每次迭代称为一个epoch。

每个epoch由两部分组成：

训练循环：迭代训练数据集，并尝试收敛到最佳参数。

验证循环:迭代测试集检查模型是否正在改进

让我们简单地熟悉一下培训循环中使用的一些概念。跳到前面，看看优化循环的完整实现。

6.3 损失函数

当面对一些训练数据时，我们未经训练的网络可能不会给出正确的答案。损失函数计算所获结果与目标值的不相似程度，整个训练过程我们想要最小化损失函数。为了计算损失，我们使用给定数据样本的输入进行预测，并将其与真实数据标签值进行比较。

通常的损失包括nn.MSELoss，nn.NLLLoss为分类。nn.CrossEntropyLoss包括nn.

我们将模型的输出结果传递给nn.CrossEntropyLoss，它将规范化logits并计算预测误差。

# Initialize the loss function
loss_fn = nn.CrossEntropyLoss()

6.4 优化

优化是在训练过程中调整模型参数减少模型误差的过程。优化算法定义了过程的执行。在这个例子中，我们使用随机梯度下降算法。全部的优化逻辑被打包在optimizer中。在这，我们使用SGD优化器。除此之外，Pytorch中也有一些其他的优化器，例如ADAM和RMSProp，在不同的模型和数据中表现的更好。

我们初始化优化器，通过输入需要被优化的模型参数，并且传入学习率参数。

optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

在这个训练循环内部，优化发生按照如下步骤：

• 调用optimizer.zero_grad()去重置模型参数的梯度。默认梯度增加，阻止重复计算，每个迭代设置为0。

• 反向传播loss通过调用loss.backward()。Pytorch存储loss对每个参数的梯度。

• 一旦我们有了梯度，我们调用optimizer.step()去调整参数。

6.5 完整实现

我们定义train_loop循环我们的优化代码，test_loop根据我们的测试数据去评估模型的表现。

def train_loop(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    for batch, (X, y) in enumerate(dataloader):
        # Compute prediction and loss
        pred = model(X)
        loss = loss_fn(pred, y)

        # Backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if batch % 100 == 0:
            loss, current = loss.item(), batch * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")


def test_loop(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    test_loss, correct = 0, 0

    with torch.no_grad():
        for X, y in dataloader:
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()

    test_loss /= num_batches
    correct /= size
    print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")

我们初始化损失函数和优化器，传递它到train_loop和test_loop。自由增长epochs去提高模型。

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

epochs = 10
for t in range(epochs):
    print(f"Epoch {t+1}\n-------------------------------")
    train_loop(train_dataloader, model, loss_fn, optimizer)
    test_loop(test_dataloader, model, loss_fn)
print("Done!")

7、保存并且加载模型

在这部分，在本节中，我们将了解如何通过保存、加载和运行模型预测来保持模型状态。

import torch
import torchvision.models as models

7.1 保存和加载模型权重

Pytorch存储学习到的参数在一个内部状态字典中，叫做state_dict。这些能被保存经由torch.save()模型。

model = models.vgg16(pretrained=True)
torch.save(model.state_dict(), 'model_weights.pth')

为了加载模型权重，你需要去创建一个相同模型的例子，然后使用load_state_dict()方法去加载参数。

model = models.vgg16() # we do not specify pretrained=True, i.e. do not load default weights
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()

注意：在推断之前一定要调用model.eval()方法，去设置dropout和批量标准化为评价模式。否则将会导致不一样的推断结果。

7.2 保存并且加载模型形状

当加载模型权重的时候，我们首先需要去初始化模型类，因为类定义了网络结构。我们也许想要和模型一起存储类的结构，我们能够传递模型并且保存函数。

torch.save(model, 'model.pth')

然后我们加载模型如下：

model = torch.load('model.pth')