Pytorch学习(三）—— 常用工具包

数据集读取

在Pytorch中，数据加载可以通过自定义的数据集对象实现。数据集对象被抽象为Dataset类，实现自定义的数据集需要继承Dataset，并实现两个Python的魔方方法：

• _ _ getitem _ _： 返回一条数据或一个样本。
• _ _ len _ _： 返回样本的数量。

以Kaggle Dog vs. Cat 数据集为例：

import os
import torch
from torch.utils import data
from PIL import Image
import numpy as np

'''
    DataSet
'''
class DogCat(data.Dataset):
    def __init__(self, root):
        imgs = os.listdir(root)
        self.imgs = [os.path.join(root, img) for img in imgs]

    def __getitem__(self, index):
        img_path = self.imgs[index]
        lable = 1 if 'dog' in img_path.split('/')[-1] else 0
        pil_img = Image.open(img_path)
        array = np.asarray(pil_img)
        data = torch.from_numpy(array)
        return data, lable
    def __len__(self):
        return len(self.imgs)

dataset = DogCat('../../../DogCat/')
img, label = dataset[0]

数据预处理和ImageFolder

数据预处理

Pytorch提供的torchvision视觉工具包提供了很多视觉图像处理的工具，其中transforms模块提供了对PIL Image对象和Tensor对象的常用操作。

对PIL Image的常见操作如下：

• Resize： 调整图片尺寸。
• Crop： 裁剪图片。
• Pad: 填充。
• ToTensor： 将PIL Image对象转成Tensor，会自动将[0, 255]归一化为[0,1]。

对Tensor的常见操作如下：

• Normalize: 标准化，即减均值，除以标准差。
• ToPILImage： 将Tensor转换为PIL Image对象。

可以使用Compose操作将多个预处理操作组合起来：

from torchvision import transforms as T

'''
    torchvision
'''

transforms = T.Compose([
    T.Resize(224),          # 缩放，保持长宽比不变，最短边为224
    T.CenterCrop(224),      # 从图片中间切出224 * 224的图片
    T.ToTensor(),           # 将图片(Image)转成Tensor，归一化至[0, 1]
    T.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5])        # 标准化至[-1, 1]
])

class PreDogCat(data.Dataset):
    def __init__(self, root):
        imgs = os.listdir(root)
        self.imgs = [os.path.join(root, img) for img in imgs]
        self.transforms = transforms

    def __getitem__(self, item):
        img_path = self.imgs[item]
        lable = 1 if 'dog' in img_path.split('/')[-1] else 0
        data = Image.open(img_path)
        if self.transforms:
            data = self.transforms(data)
        return data, lable

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

dataset = PreDogCat('../../../DogCat/')

for img, label in dataset:
	print(img.size(), label)

除此之外，transforms还可以通过Lambda封装自定义的转换策略：

trans = T.Lambda(lambda img: img.rotate(random() * 360))

ImageFolder

torchvision中预先实现了常用的Dataset，包括CIFAR-10，ImageNet，COCO，MNIST、LSUN等数据集，可以通过调用torchvision.datasets下相应对象来调用相关数据集。

此外，torchvision中还定义了一个常用的Dataset——ImageFolder，它假设所有文件按照文件夹保存，每个文件夹下存储同一个类别的图片，文件夹名为类名，其构造函数如下：

    def __init__(self, root, transform=None, target_transform=None,
                 loader=default_loader, is_valid_file=None):

主要参数如下：

• root： 在指定root路径下寻找图片。
• transform： 对PIL Image进行的转换操作，transform的输入是使用loader函数读取的图片对象。
• target_transform： 对label的转换。
• loader： 指定加载图片的函数，默认操作是读取为PIL Image对象。

from torchvision.datasets import ImageFolder

# 不加transform的ImageFolder
dataset = ImageFolder('../../../DogCat/')
print(dataset.class_to_idx)
print(dataset.imgs)
dataset[0][0].show()        # image
print(dataset[0][1])        # label

# 添加transform的ImageFolder
normalize = T.Normalize(mean=[0.4, 0.4, 0.4], std=[0.2, 0.2, 0.2])
transforms = T.Compose([
    T.RandomResizedCrop(224),
    T.RandomHorizontalFlip(),
    T.ToTensor(),
    normalize
])

dataset = ImageFolder('../../../DogCat/', transform=transforms)

toimg = T.ToPILImage()
toimg(dataset[0][0] * 0.2 + 0.4)

数据加载

在训练神经网络时，通常都是对一个batch的数据进行操作，同时还需要对数据进行shuffle和并行加速。对此，Pytorch提供了DataLoader进行数据加载。

DataLoader的函数定义如下：

    def __init__(self, dataset, batch_size=1, shuffle=False, sampler=None,
                 batch_sampler=None, num_workers=0, collate_fn=None,
                 pin_memory=False, drop_last=False, timeout=0,
                 worker_init_fn=None, multiprocessing_context=None):

• dataset: 加载的数据集，Dataset对象。
• batch_size: 批大小。
• shuffle： 是否将数据打乱。
• sampler： 样本抽取方式。
• num_worker: 使用多线程加载的进程数。
• collate_fn: 如何将多个样本拼接成一个batch。
• pin_memory： 是否将数据保存在pin memory区，pin memory中的数据转移到GPU中会快一些。
• drop_last： dataset中的数据个数可能不是batch_size的整数倍，drop_last为True时会将多出来不足一个batch的数据丢弃。

另外，Pytorch中单独提供了一个sampler模块，用来对数据进行采样。当DataLoader的shuffle参数为Ture时，系统会自动调用随机采样器RandomSampler，否则调用默认采样器SequentialSampler，此外还提供了另一个很有用的采样器：WeightRandomSampler，它会根据每个样本的权重选取数据。

from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
from torchvision import transforms as T
from torch.utils.data.sampler import WeightedRandomSampler

transforms = T.Compose([
    T.RandomResizedCrop(224),
    T.ToTensor()
])
dataset = ImageFolder("../../../DogCat", transform=transforms)
dataloader = DataLoader(dataset, batch_size=20, shuffle=True, num_workers=0, drop_last=False)


for batch_datas, batch_labels in dataloader:
    print(batch_datas.size(), batch_labels.size())

weights = [2 if lable == 1 else 1 for data, lable in dataset]

sampler = WeightedRandomSampler(weights, num_samples=10, replacement=True)

dataloader = DataLoader(dataset, batch_size=5, sampler=sampler)

当DataLoader使用多进程加载时，程序出现异常终止，比如Ctrl+C强制退出时，相应的数据加载进程可能无法正常退出，出现程序已经退出但是GPU显存和内存依旧被占用的情况，这时候就需要使用如下命令手动强行终止进程：
ps x | grep <cmdline> | awk '{print $1}' | xargs kill

• ps x： 获取当前用户的所有进程。
• grep < cmdline >： 找到已经停止的Pytorch程序进程。
• awk ‘{print $1}’： 获取进程pid。
• xargs kill： kill进程。

在kill之前，建议先打印确认进程：
ps x | grep <cmdline>

torchvision

除了上述的datasets数据集和transforms数据预处理操作，torchvision中还提供了models。

models中提供了深度学习中各种经典网络的网络结构及预训练模型：

from .alexnet import *
from .resnet import *
from .vgg import *
from .squeezenet import *
from .inception import *
from .densenet import *
from .googlenet import *
from .mobilenet import *
from .mnasnet import *
from .shufflenetv2 import *

加载模型：

from torchvision import models
from torch import nn

'''
    torchvision models
'''

resnet34 = models.resnet34(pretrained=True, num_classes=1000)
# resnet34.fc = nn.Linear(512, 10)        #   修改为10分类问题

# 模型保存
torch.save(resnet34.state_dict(), 'resnet34.pth')

# 模型加载
resnet34.load_state_dict(torch.load('resnet34.pth'))

# 优化器参数保存
optimizer = torch.optim.Adam(resnet34.parameters(), lr=0.01)
torch.save(optimizer.state_dict(), 'optimizer.pth')

# 优化器参数加载
optimizer.load_state_dict(torch.load('optimizer.pth'))

all_data = dict(
    optimizer=optimizer.state_dict(),
    model=resnet34.state_dict(),
    info='模型和优化器的所有参数'
)

torch.save(all_data, 'all.pth')

all_data = torch.load('all.pth')
all_data.key()

torchvision还提供了两个常用函数：一个make_grid，它能将多张图片拼接在一个网格中；另一个是save_img，它能将Tensor保存为图片。

GPU加速：CUDA

Pytorch中的Tensor和nn.Module都带有一个.cuda方法，调用此方法可以将其转换为对应的GPU对象。需要注意的是，tensor.cuda会返回一个新对象，这个新对象的数据已转移到GPU，而之前的tensor数据还在原设备上。module.cuda会将所有的数据都迁移到GPU，并返回自己，所以 module = module.cuda 与 module.cuda效果相同。

为什么将数据转移至GPU的方法叫做.cuda而不叫.gpu?
这是因为GPU的编程接口采用CUDA，而目前并不是所有的GPU都支持CUDA，只有部分NVIDIA的GPU才支持。AMD GPU的编程接口采用OpenCL，因此Pytorch中预留了.cl方法，用于支持AMD GPU。

关于使用GPU加速的建议：

• GPU运算很快，但是当计算量小时，并不能体现出它的优势，因此一些简单的操作可以直接利用CPU完成。
• 数据在CPU和GPU之间传递比较耗时，应当尽量避免。
• 在低精度计算时，可以考虑HalfTensor，相比FloatTensor能够节省一半的显存，但是需要注意数值溢出的问题。
• Pytorch封装了分布式训练接口。