pytorch入门（三）

在前两篇文章中我们了解了pytorch的一些基础概念以及模型训练的整体框架。今天我们继续介绍pytorch的一些常用包。

传送门：

pytorch入门（一）

pytorch入门（二）

Dataset&DataLoader

对于数据处理上，我们希望样本数据集代码和模型训练代码分开，这样便于获得更强的可读性和模块化。因此，pytorch提供了两个数据元素：

from torch.utils.data import Dataset
from torch.utils.data import DataLoader

Dataset存储样本示例以及其对应的标签；DataLoader包装了数据集的可迭代对象，方便访问样本，可以认为是一个数据加载器。

Dataset是一种抽象类，所有继承它的子类都应该包含__len__方法和__getitem__方法，前者表示数据集大小，后者支持实例化后的样本整数索引。一般情况下，继承此类的对象也都包含初始化方法__init__.

下面举一个例子：

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

img_dir是存储数据(这个例子应该是图片)路径，annotations_file是存储图片对应标签的CSV文件。然后，方法__len__返回的是数据集所有图片的个数；__getitem__方法是在给定样本的索引后，加载样本数据以及其对应的标签。两个tranform是做一些必要的转换，这里先不用管。

在训练模型的时候，一般只取批量样本进行训练，因此，DataLoader实现这种“小批量”训练过程：

from torch.utils.data import DataLoader


DataLoader(dataset, batch_size=1, shuffle=False, 
           sampler=None, num_workers=0,  drop_last=False)

dataset就是要加载的数据集；batch_size表示一个批量加载多少个样本；shuffle表示在每个epoch是否重新打乱数据；sampler定义从数据集提取样本的策略，如果给定，则shuffle参数不生效；num_workers表示用多少个子进程加载数据，0表示将在主进程加载；drop_last表示如果数据集大小不能被batch_size整除，是否删除最后一个不完整的batch，默认是不删除。

DataLoader返回的是一个迭代器，说明可以用next()或者yield方法接收。而在pytorch中，Dataset一般封装在DataLoader里面，因此，DataLoader返回的迭代器是指向__getitem__方法的返回的东西，比如上面的例子返回的是image和label：

dataloader = DataLoader()
for image,label in dataloader:
    print(image,labels)

torchvision

torchvision是pytorch的一个包，它包含了流行的数据集、模型结构以及常用的图片转换工具。它包含以下四块内容：

• torchvision.datasets

• torchvison.transforms

• torchvision.models

• torchvision.utils
  
  

torchvision.datasets包含了常用的数据集，下面列出了一些数据集的加载方法：

from torchvision import datasets


datasets.MNIST()
datasets.CocoCaptions()
datasets.LSUN()
datasets.ImageFolder()
datasets.CIFAR10()
datasets.CIFAR100()
datasets.STL10()

所有的数据集都是继承了torch.utils.data.Dataset类方法，因此，它们都具有__len__和__getitem__方法，也都可以用采集器DataLoader加载样本。除此之外，它们都有两个共同的参数——tansform和target_transfotm。现在通过MNIST数据集进行介绍。

MNIST(root,train=True,transform=None,target_transform=None,download=False)

• root:存储样本的主目录

• train:True为训练集,False为测试机

• download:是否从互联网下载数据集
  

• transform:将原始图片做转换的函数

• target_transform:输入为样本标签,输出是对其做转换
  

最后两个参数是转换器，本质是函数，它传入的参数就是torchvision第二个模块的内容——transforms，我们可以打印以下一共有多少种转换函数：

我们就简单的介绍几个.

• torchvision.transforms.ToTensor

把一个取值范围是[0,255]的PIL.Image或者shape为(H,W,C)的numpy.ndarray，转换为（C,H,W)的torch.FloadTensor

data = np.random.randint(0, 255, size=300)
img = data.reshape(10,10,3)
print(img.shape)
img_tensor = transforms.ToTensor()(img) # 转换成tensor
print(img_tensor)

• torchvision.transforms.ToPILImage

将形状为（C,H,W)的Tensor或ndarray转换为PIL.Image

•  torchvision.transforms.Compose

封装多个transform的函数，即转化器组合，下面举个例子：

transforms.Compose([
     transforms.CenterCrop(10),
    transforms.ToTensor(),
])

其余更多的转换器见官网，这里不再赘述了。

torchvision.models包含了一些常用的模型结构，这个等到介绍迁移学习的时候再和tensorflow的一起介绍，这里略过。然后剩下最后一个模块——torchvision.utils，它是一个制作图像网络的模块：

torchvision.utils.make_grid(
       tensor, nrow=8, padding=2, normalize=False, 
       range=None, scale_each=False, pad_value=0)

tensor是四维度的张量(B,C,H,W)或者形状相同的图像列表；nrows表示网格中的行数，最终网格大小B/nrow；normalize是否归一化图像;range表示归一化后的最大最小值;scale_each表示是否在所有图像上归一化还是仅在自身上归一化;pad_value表示填充像素的取值;padding是填充的步数。

利用cuda训练模型

pytorch支持使用GPU训练模型，大幅度提高训练速率，通过下面代码查看计算机是否有可用的cuda：

torch.cuda.is_available()

下面通过一个简单的神经网络来看看如何使用cuda进行模型训练。

首先，定义一个DNN模型：

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

然后，cuda或者cpu作为device传入模型：

device = 'cuda' if torch.cuda.is_available() else 'cpu'
print('Using {} device'.format(device))


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

这里不仅仅可以把模型复制到GPU上，也可以把任何张量复制到GPU:

mytensor = my_tensor.to(device)

那么如何利用GPU进行模型多核训练呢，利用torch.nn的DataParallel实现核心代码：

model = Model(input_size, output_size)
if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUs!")
    model = nn.DataParallel(model)


model.to(device)

首先我们需要实例化一个模型，然后验证我们是否有多个 GPU。如果有，那么用 nn.DataParallel来包裹模型。然后用 model.to(device)把模型放到多 GPU 中。

保存与加载模型

pytorch保存和加载模型使用三个核心功能：

• torch.save：将pickle序列化对象保存到磁盘

• torch.load：将pickle文件反序列化到内存

• torch.nn.Module.load_state_dict：用函数state_dict加载模型参数字典

state_dict被称为状态字典，它将神经网络的每一层映射到其参数中，key是参数名称（权重或者偏移项），value是取值。举个例子，下面是模型和优化器的参数：

# 打印模型的状态字典
print("Model's state_dict:")
for param_tensor in model.state_dict():
    print(param_tensor, "\t", model.state_dict()[param_tensor].size())
# 打印优化器的状态字典
print("Optimizer's state_dict:")
for var_name in optimizer.state_dict():
    print(var_name, "\t", optimizer.state_dict()[var_name])

加载和保存完整模型最简单的方式可用如下代码：

torch.save(model, PATH)


model = torch.load(PATH)
model.eval()

注意，在运行推理前，必须调用model.eval()设置 dropout 和 batch normalization 层为评估模式。如果不这么做，可能导致模型推断结果不一致。

官方推荐使用的是带状态字典的保存/加载方式：

torch.save(model.state_dict(), PATH)


model = TheModelClass(*args, **kwargs)  #训练的模型
model.load_state_dict(torch.load(PATH))
model.eval()

将模型保存/加载到GPU可使用如下方式：

torch.save(model.state_dict(), PATH)


device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device) # 确保在你提供给模型的任何输入张量上调用input = input.to(device)

将模型保存到CPU，加载到GPU的方式如下：

torch.save(model.state_dict(), PATH)


device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
# Choosewhatever GPU device number you want
model.load_state_dict(torch.load(PATH, map_location="cuda:0")) 
model.to(device) 
# 确保在你提供给模型的任何输入张量上调用input = input.to(device)

将在CPU上训练的模型加载到GPU时，需要指定GPU设备，这在torch.load()函数的map_location参数设置：cuda:device_id.

当模型在GPU上时，传入的张量一定要声明为cuda优化器：

my_tensor = my_tensor.to(torch.device('cuda'))

如果模型训练过程中断了，并在下一次想继续训练或者用于推断，则保存和加载方式如下：

torch.save({
     'epoch': epoch,
     'model_state_dict': model.state_dict(),
     'optimizer_state_dict': optimizer.state_dict(),
     'loss': loss,
     ...
    }, PATH)

model = TheModelClass(*args, **kwargs)
optimizer = TheOptimizerClass(*args, **kwargs)
checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']
# 用于推断
model.eval()
# 用于再训练
model.train()

对于DataParallel模型，保存方式如下：

torch.save(model.module.state_dict(), PATH)

在迁移学习中，要加载部分模型，可以使用如下方式：

modelB = TheModelBClass(*args, **kwargs)
modelB.load_state_dict(torch.load(PATH), strict=False)

查看模型训练进度

最后，我们介绍一个python三方库，用于查看模型训练的进度——tqdm。下面通过一个例子来说明：

from tqdm import tqdm
from tqdm.notebook import tqdm_notebook #进度条


num_epochs = 100
get_num_batches = 10


epoch_bar = tqdm_notebook(desc='training routine', total=num_epochs,position=0)
train_bar = tqdm_notebook(desc='split=train',total=num_epochs*get_num_batches, position=1, leave=True)
val_bar = tqdm_notebook(desc='split=val',total=num_epochs*get_num_batches*0.5, position=1, leave=True)


for epoch_index in range(num_epochs):
    running_loss = 1
    running_acc = 0.9
    for i in range(get_num_batches):
        train_bar.set_postfix(loss=running_loss, acc=running_acc, epoch=epoch_index)
        train_bar.update()
    for i in range(int(get_num_batches*0.5)):
        train_bar.set_postfix(loss=running_loss, acc=running_acc, epoch=epoch_index)
        val_bar.update()


    epoch_bar.update()

我们只说明tqdm_notebook的一些参数即可：

iterable:可迭代的对象

desc:进度条的描述，可以理解为命名

total:一共有多少次进度

position:进度条的位置，从0开始

leave:进度完成后是否保留进度条

这里不立刻传入迭代对象，而是在每一次迭代后通过update()方法更新，其中set_postfix是指给进度变添加描述，只不过是在后面。我们查看动态效果图：

好了，pytorch入门就介绍到这。下次我们用一个例子把这些所学结合起来，看看一个项目如何从获取数据到模型落地的。

参考资料：

https://www.ptorch.com/docs/1/utils-data

https://pytorch-cn.readthedocs.io/zh/latest/torchvision/torchvision-datasets/

https://www.ptorch.com/docs/1/transforms