DATASETS & DATALOADERS
两个有用的,torch.utils.data.DataLoader
and torch.utils.data.Dataset
,其中Dataset
存储样本和标签(就是图片和真值),而DataLoader
把它变成可迭代对象。
加载数据集
例子是加载 Fashion-MNIST数据集
import torch
from torch.utils.data import Dataset
from torchvision import datasets
from torchvision.transforms import ToTensor, Lambda
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() )
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()
后面是可视化的部分,不研究它,主要研究那个加载数据集的函数参数。
其中root
是你放训练集/测试集的路径;train
是指定它是训练集还是测试集;download
是看你要不要下载这个数据集;transform
是指定要对数据做什么变换。
其中的ToTensor()
代表用这个函数来处理载入的数据。ToTensor()
函数其接受PIL Image或numpy.ndarray格式,功能是:先把HWC改为CHW格式;在转换为float类型;在除以255转为0-1之间。
用自己的文件定制一个数据集
一个定制的数据集的类必须有一下三个方法: __init__
, __len__
, and __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)
sample = {"image": image, "label": label}
return sample
重写了这三个方法,就是你自己定义了一个数据集的class了。里面的target_transform
可以很明显看出来是对标签进行处理的。__init__
方法初始化实例,传进去四个:图片路径,标签路径,两个变换的函数。 __len__
返回图片样本的数量;__getitem__
方法是:你给一个idx
,返回一个样本的dict
,其中加载了两种数据并做了变换。
这样你就定义了一个自己的数据集的class。
用DataLoaders准备你的训练数据
用Dataset
一次返回一个样本和标签,但是训练时希望:输送minibatches
而不是一次一个;希望每个epoch进行reshuffle
重新洗牌以减少过拟合;希望可以用python的multiprocessing
来加速取数据。所以有了我们的DataLoader
。
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)
# 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}")
把数据集,batchsize,是否洗牌等参数传进去以后,就返回了一个DataLoader
的实例;这个实例是一个可迭代对象,通过iter()
函数变为迭代器,通过next()
来调取一个元素,运行代码可以看出一次拿出来64张,
TRANSFORMS
-transform
to modify the features and target_transform
to modify the labels .看看官网的例子:
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))
)
里面的部分刚刚都说了,但是这个target_transform
里面的这个Lambda()
我就没搞懂,看了一下后边才知道,这是torchvision.transforms
里面提供的一个功能,把用户定义的lambda函数传进去,做个判断,看能不能调用,下面是这个函数的代码,作用不大,不纠结,会用就成。
class Lambda(object):
"""Apply a user-defined lambda as a transform.
Args:
lambd (function): Lambda/function to be used for transform.
"""
def __init__(self, lambd):
assert callable(lambd), repr(type(lambd).__name__) + " object is not callable"
self.lambd = lambd
def __call__(self, img):
return self.lambd(img)
def __repr__(self):
return self.__class__.__name__ + '()'
BUILD THE NEURAL NETWORK
神经网络包含很多对数据进行运算的layers/modules,torch.nn
这个namespace提供了所有的这些操作。
import os
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
#选择运算的设备
device = 'cuda' if torch.cuda.is_available() else 'cpu'
#继承torch.nn,建立一个自己的类,重写init和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),
nn.ReLU()
)
def forward(self, x):
x = self.flatten(x)
logits = self.linear_relu_stack(x)
return logits
#创建一个神经网络的实例,放到device中。
model = NeuralNetwork().to(device)
#传入input数据,就开始了forward,不要直接调用forward。
X = torch.rand(1, 28, 28, device=device)
logits = model(X)
pred_probab = nn.Softmax(dim=1)(logits)
y_pred = pred_probab.argmax(1)
以上就是建立模型,输入数据,运算得出结果的部分。模型计算出来的结果就是全连接层最后的结果,是无界的,可正可负,代码里叫它logits
,然后用Softmax
进行归一化,维度是1,就是按照行进行归一化,意思是每一行的数值都归一化到0-1之间,并且相加为1,就是得出了每一类的概率。
关于Softmax
,官方教程上这样说,很好理解:
The last linear layer of the neural network returns logits - raw values in [-infty, infty] - which are passed to the nn.Softmax module. The logits are scaled to values [0, 1] representing the model’s predicted probabilities for each class. dim parameter indicates the dimension along which the values must sum to 1.
然后用argmax
函数找出每一行最大值的索引值(维度是1,所以是按照每一行找最大值)
然后看nn.Flatten
函数,源码里这个函数的注释如下,一看就懂:
Flattens a contiguous range of dims into a tensor. For use with :class:`~nn.Sequential`.
Shape:
- Input: :math:`(N, *dims)`
- Output: :math:`(N, \prod *dims)` (for the default case).
Args:
start_dim: first dim to flatten (default = 1).
end_dim: last dim to flatten (default = -1).
就是它这个函数一般跟着nn.Sequential
一起用,除了第一个维度不变,剩下的拉平,当然,你也可以指定拉平的开始和结束的维度,默认是1和-1。
要快速构建网络,用nn.Sequential
,就是按照你定义的模块顺序来执行计算的顺序,不用forward函数指定运算顺序。
seq_modules = nn.Sequential(
flatten,
layer1,
nn.ReLU(),
nn.Linear(20, 10)
)
input_image = torch.rand(3,28,28)
logits = seq_modules(input_image)
我个人理解是nn.Sequential
和nn.Module
可以混动,那个方便用哪个,就是说nn.Module
里面可以嵌套nn.Sequential
,把它当成一个像普通卷积操作一样的来用,反正它也是按照顺序来的。反之应该也可以(没有实践,只是我的理解)。
怎么知道你定义的这个模型的所有参数呢?用Model Parameters
。只要你继承了nn.Module
建立类和实例,所有参数自动就有了,可以用模型的parameters()
or named_parameters()
方法来访问参数。
torch.autograd
训练模型时,常用反向传播,就是根据损失函数求出对应参数的梯度,torch.autograd
就是根据计算图(computational graph)自动计算梯度的。贴上官方教程代码例子:
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)
print('Gradient function for z =',z.grad_fn)
print('Gradient function for loss =', loss.grad_fn)
loss.backward()
print(w.grad)
print(b.grad)
要求梯度的参数可以在创建的时候加requires_grad=True
,也可以以后用x.requires_grad_(True)
方法来加上这个属性。
调用了loss.backward()
,反向传播就开始了,就是计算各个参数的梯度。计算谁的梯度呢?
requires_grad=True
的,你想,变量x作为input为 不计算梯度,就是因为没有设定True
,一开始我也很迷糊,现在很清晰了,就是说哪个是参数,就设为True
,其他的就是数据。pytorch就是通过这个参数来识别哪个要优化,哪个不优化。- 一般是只有计算图中设为
True
,且时叶子节点的,才有梯度。 - 一个计算图一般只能
backward
一次,为了节约资源,算完了就释放了,如果你想在同样的计算图中backward
几次,调用backward
的时候设定参数retain_graph=True
。
不想后向传播的时候,用torch.no_grad()
和detach()
。
z = torch.matmul(x, w)+b
print(z.requires_grad)
with torch.no_grad():
z = torch.matmul(x, w)+b
print(z.requires_grad)
z = torch.matmul(x, w)+b
z_det = z.detach()
print(z_det.requires_grad)
使用场景有二:
- 只是检测,不训练,不用后向传播
- 只是微调某些参数,让另一部分不调整。
OPTIMIZING MODEL PARAMETERS优化模型参数
超参Hyperparameters:
- Number of Epochs - the number times to iterate over the dataset。一次迭代叫一个epoch。
- Batch Size - the number of data samples seen by the model in each epoch。每次迭代处理的数据样本的个数叫Batch Size。
- Learning Rate - how much to update models parameters at each batch/epoch. Smaller values yield slow learning speed, while large values may result in unpredictable behavior during training.每次更新参数的时候更新的步长叫
Learning Rate
。太小了学得慢,太大了不稳定。
在优化的循环Optimization Loop中,每个epoch包含2个主要部分:
- The Train Loop - iterate over the training dataset and try to converge to optimal parameters.训练循环,在训练数据上迭代以便收敛。
- The Validation/Test Loop - iterate over the test dataset to check if model performance is improving.验证循环,在验证数据上迭代,来检验模型的表现是否在提升。
Optimizer优化器:
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
定义一个优化器,传入的参数是模型里要优化的参数和超参:学习率。
在训练循环 training loop 中,优化有三步:
- 调用
optimizer.zero_grad()
来使梯度清零,因为梯度是累积的,为了防止重复计算,每次迭代以后都要清零。 - 后向传播,计算梯度,
loss.backwards()
- 有了梯度,然后优化参数,
optimizer.step()
SAVE AND LOAD THE MODEL保存/加载模型
pytorch模型把学习到的参数存储在内部的一个状态字典(an internal state dictionary)里,叫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.train()
,使模型处于这个模式之下,此时模型会启用dropout and batch normalization
,如果是验证/测试,那么要先将模型设置model.eval()
,不启用刚刚这两个功能。
使用model.eval()
还是会计算梯度的,只是不用那两个功能而已,但是用with torch.no_grad()
,可以停止autograd
模块的工作,可以加速,节省显存,不更新梯度了,但是这个并不影响刚刚提到的那两个功能。
刚刚仅仅是保存了模型的可学习参数,但是模型结构并没有保存。yaoDATASETS & DATALOADERS
两个有用的,torch.utils.data.DataLoader
and torch.utils.data.Dataset
,其中Dataset
存储样本和标签(就是图片和真值),而DataLoader
把它变成可迭代对象。加载数据集例子是加载 Fashion-MNIST数据集pythonimport torchfrom torch.utils.data import Datasetfrom torchvision import datasetsfrom torchvision.transforms import ToTensor, Lambdaimport matplotlib.pyplot as plttraining_data = datasets.FashionMNIST(root="data",train=True,download=True,transform=ToTensor() )test_data = datasets.FashionMNIST(root="data",train=False,download=True,transform=ToTensor() )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, 3for 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()
后面是可视化的部分,不研究它,主要研究那个加载数据集的函数参数。其中root
是你放训练集/测试集的路径;train
是指定它是训练集还是测试集;download
是看你要不要下载这个数据集;transform
是指定要对数据做什么变换。其中的ToTensor()
代表用这个函数来处理载入的数据。ToTensor()
函数其接受PIL Image或numpy.ndarray格式,功能是:先把HWC改为CHW格式;在转换为float类型;在除以255转为0-1之间。用自己的文件定制一个数据集一个定制的数据集的类必须有一下三个方法: __init__
, __len__
, and __getitem__
。比如FashionMNIST 数据集,假设它的图片存在img_dir
里,标签的CSV文件存在annotations_file
里面。下面看代码:pythonimport osimport pandas as pdfrom torchvision.io import read_imageclass 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) sample = {"image": image, "label": label} return sample
重写了这三个方法,就是你自己定义了一个数据集的class了。里面的target_transform
可以很明显看出来是对标签进行处理的。__init__
方法初始化实例,传进去四个:图片路径,标签路径,两个变换的函数。 __len__
返回图片样本的数量;__getitem__
方法是:你给一个idx
,返回一个样本的dict
,其中加载了两种数据并做了变换。这样你就定义了一个自己的数据集的class。用DataLoaders准备你的训练数据用Dataset
一次返回一个样本和标签,但是训练时希望:输送minibatches
而不是一次一个;希望每个epoch进行reshuffle
重新洗牌以减少过拟合;希望可以用python的multiprocessing
来加速取数据。所以有了我们的DataLoader
。pythonfrom torch.utils.data import DataLoadertrain_dataloader = DataLoader(training_data, batch_size=64, shuffle=True)test_dataloader = DataLoader(test_data, batch_size=64, 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}")
把数据集,batchsize,是否洗牌等参数传进去以后,就返回了一个DataLoader
的实例;这个实例是一个可迭代对象,通过iter()
函数变为迭代器,通过next()
来调取一个元素,运行代码可以看出一次拿出来64张,TRANSFORMS-transform
to modify the features and target_transform
to modify the labels .看看官网的例子:pythonfrom torchvision import datasetsfrom torchvision.transforms import ToTensor, Lambdads = 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)))
里面的部分刚刚都说了,但是这个target_transform
里面的这个Lambda()
我就没搞懂,看了一下后边才知道,这是torchvision.transforms
里面提供的一个功能,把用户定义的lambda函数传进去,做个判断,看能不能调用,下面是这个函数的代码,作用不大,不纠结,会用就成。pythonclass Lambda(object): """Apply a user-defined lambda as a transform. Args: lambd (function): Lambda/function to be used for transform. """ def __init__(self, lambd): assert callable(lambd), repr(type(lambd).__name__) + " object is not callable" self.lambd = lambd def __call__(self, img): return self.lambd(img) def __repr__(self): return self.__class__.__name__ + '()'
BUILD THE NEURAL NETWORK神经网络包含很多对数据进行运算的layers/modules,torch.nn
这个namespace提供了所有的这些操作。pythonimport osimport torchfrom torch import nnfrom torch.utils.data import DataLoaderfrom torchvision import datasets, transforms#选择运算的设备device = 'cuda' if torch.cuda.is_available() else 'cpu'#继承torch.nn,建立一个自己的类,重写init和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), nn.ReLU() ) def forward(self, x): x = self.flatten(x) logits = self.linear_relu_stack(x) return logits#创建一个神经网络的实例,放到device中。model = NeuralNetwork().to(device)#传入input数据,就开始了forward,不要直接调用forward。X = torch.rand(1, 28, 28, device=device)logits = model(X)pred_probab = nn.Softmax(dim=1)(logits)y_pred = pred_probab.argmax(1)
以上就是建立模型,输入数据,运算得出结果的部分。模型计算出来的结果就是全连接层最后的结果,是无界的,可正可负,代码里叫它logits
,然后用Softmax
进行归一化,维度是1,就是按照行进行归一化,意思是每一行的数值都归一化到0-1之间,并且相加为1,就是得出了每一类的概率。关于Softmax
,官方教程上这样说,很好理解:*The last linear layer of the neural network returns logits - raw values in [-infty, infty] - which are passed to the nn.Softmax module. The logits are scaled to values [0, 1] representing the model’s predicted probabilities for each class. dim parameter indicates the dimension along which the values must sum to 1.*然后用argmax
函数找出每一行最大值的索引值(维度是1,所以是按照每一行找最大值)然后看nn.Flatten
函数,源码里这个函数的注释如下,一看就懂:python Flattens a contiguous range of dims into a tensor. For use with :class:`~nn.Sequential`. Shape: - Input: :math:`(N, *dims)` - Output: :math:`(N, \prod *dims)` (for the default case). Args: start_dim: first dim to flatten (default = 1). end_dim: last dim to flatten (default = -1).
就是它这个函数一般跟着nn.Sequential
一起用,除了第一个维度不变,剩下的拉平,当然,你也可以指定拉平的开始和结束的维度,默认是1和-1。要快速构建网络,用nn.Sequential
,就是按照你定义的模块顺序来执行计算的顺序,不用forward函数指定运算顺序。pythonseq_modules = nn.Sequential( flatten, layer1, nn.ReLU(), nn.Linear(20, 10))input_image = torch.rand(3,28,28)logits = seq_modules(input_image)
我个人理解是nn.Sequential
和nn.Module
可以混动,那个方便用哪个,就是说nn.Module
里面可以嵌套nn.Sequential
,把它当成一个像普通卷积操作一样的来用,反正它也是按照顺序来的。反之应该也可以(没有实践,只是我的理解)。怎么知道你定义的这个模型的所有参数呢?用Model Parameters
。只要你继承了nn.Module
建立类和实例,所有参数自动就有了,可以用模型的parameters()
or named_parameters()
方法来访问参数。torch.autograd训练模型时,常用反向传播,就是根据损失函数求出对应参数的梯度,torch.autograd
就是根据计算图(computational graph)自动计算梯度的。贴上官方教程代码例子:pythonimport torchx = torch.ones(5) # input tensory = torch.zeros(3) # expected outputw = torch.randn(5, 3, requires_grad=True)b = torch.randn(3, requires_grad=True)z = torch.matmul(x, w)+bloss = torch.nn.functional.binary_cross_entropy_with_logits(z, y)print('Gradient function for z =',z.grad_fn)print('Gradient function for loss =', loss.grad_fn)loss.backward()print(w.grad)print(b.grad)
要求梯度的参数可以在创建的时候加requires_grad=True
,也可以以后用x.requires_grad_(True)
方法来加上这个属性。调用了loss.backward()
,反向传播就开始了,就是计算各个参数的梯度。计算谁的梯度呢? 1. requires_grad=True
的,你想,变量x作为input为 不计算梯度,就是因为没有设定True
,一开始我也很迷糊,现在很清晰了,就是说哪个是参数,就设为True
,其他的就是数据。pytorch就是通过这个参数来识别哪个要优化,哪个不优化。 2. 一般是只有计算图中设为True
,且时叶子节点的,才有梯度。 3. 一个计算图一般只能backward
一次,为了节约资源,算完了就释放了,如果你想在同样的计算图中backward
几次,调用backward
的时候设定参数retain_graph=True
。不想后向传播的时候,用torch.no_grad()
和detach()
。pythonz = torch.matmul(x, w)+bprint(z.requires_grad)with torch.no_grad(): z = torch.matmul(x, w)+bprint(z.requires_grad)``````pythonz = torch.matmul(x, w)+bz_det = z.detach()print(z_det.requires_grad)
使用场景有二: - 只是检测,不训练,不用后向传播 - 只是微调某些参数,让另一部分不调整。OPTIMIZING MODEL PARAMETERS优化模型参数超参Hyperparameters: - Number of Epochs - the number times to iterate over the dataset。一次迭代叫一个epoch。 - Batch Size - the number of data samples seen by the model in each epoch。每次迭代处理的数据样本的个数叫Batch Size。 - Learning Rate - how much to update models parameters at each batch/epoch. Smaller values yield slow learning speed, while large values may result in unpredictable behavior during training.每次更新参数的时候更新的步长叫Learning Rate
。太小了学得慢,太大了不稳定。 在优化的循环Optimization Loop中,每个epoch包含2个主要部分: - The Train Loop - iterate over the training dataset and try to converge to optimal parameters.训练循环,在训练数据上迭代以便收敛。 - The Validation/Test Loop - iterate over the test dataset to check if model performance is improving.验证循环,在验证数据上迭代,来检验模型的表现是否在提升。Optimizer优化器:pythonoptimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
定义一个优化器,传入的参数是模型里要优化的参数和超参:学习率。在训练循环 training loop 中,优化有三步: 1. 调用optimizer.zero_grad()
来使梯度清零,因为梯度是累积的,为了防止重复计算,每次迭代以后都要清零。 2. 后向传播,计算梯度,loss.backwards()
3. 有了梯度,然后优化参数,optimizer.step()
SAVE AND LOAD THE MODEL保存/加载模型pytorch模型把学习到的参数存储在内部的一个状态字典(an internal state dictionary)里,叫state_dict
,这些参数可以用torch.save()
方法来保存。pythonmodel = models.vgg16(pretrained=True)torch.save(model.state_dict(), 'model_weights.pth')
加载参数要先创建一个模型的实例,然后用load_state_dict()
来加载参数。pythonmodel = models.vgg16() # we do not specify pretrained=True, i.e. do not load default weightsmodel.load_state_dict(torch.load('model_weights.pth'))model.eval()
训练时,模型设置model.train()
,使模型处于这个模式之下,此时模型会启用dropout and batch normalization
,如果是验证/测试,那么要先将模型设置model.eval()
,不启用刚刚这两个功能。使用model.eval()
还是会计算梯度的,只是不用那两个功能而已,但是用with torch.no_grad()
,可以停止autograd
模块的工作,可以加速,节省显存,不更新梯度了,但是这个并不影响刚刚提到的那两个功能。刚刚仅仅是保存了模型的可学习参数,但是模型结构并没有保存。要保存模型结构这个类,直接把model传入save
函数就好了。
torch.save(model, 'model.pth')
model = torch.load('model.pth')
这种方法是用python的pickle
模块实现的序列化。
Exporting Model to ONNX
pytorch是动态计算图,转为onnx时,需要传入一个size正确的测试变量。
input_image = torch.zeros((1,3,224,224))
onnx.export(model, input_image, 'model.onnx')