Pytorch入门上 —— Dataset、Tensorboard、Transforms、Dataloader

本节内容参照小土堆的pytorch入门视频教程。学习时建议多读源码，通过源码中的注释可以快速弄清楚类或函数的作用以及输入输出类型。

Dataset

借用Dataset可以快速访问深度学习需要的数据，例如我们需要访问如下训练数据：

其中，train中存放的是训练数据集，ants和bees既是文件夹名称也是其包含的图片数据的标签，val中存放的是验证数据集。

假如我们希望自己的Dataset类可以实现如下数据访问形式：

dataset = MyDataset("root_dir", "label_dir")
img, label = dataset[0]  # 通过下标访问

我们只需要继承Dataset类并覆盖其中的__getitem__()（必须）和__len__()（建议）方法。详情可在交互模式中执行如下语句查看：

from torch.utils.data import Dataset
help(Dataset)
# 或执行 Dataset??
# 或在pycharm中按住ctrl并点击Dataset直接查看源码

部分输出如下：

An abstract class representing a :class:`Dataset`.
    All datasets that represent a map from keys to data samples should subclass it. All subclasses should overwrite :meth:`__getitem__`, supporting fetching a data sample for a given key. Subclasses could also optionally overwrite :meth:`__len__`, which is expected to return the size of the dataset by many :class:`~torch.utils.data.Sampler` implementations and the default options of :class:`~torch.utils.data.DataLoader`.

定义MyDataset类：

import os

from torch.utils.data import Dataset
from torch.utils.data.dataset import T_co
from PIL import Image

class MyDataset(Dataset):
    def __init__(self, root_dir, label_dir):
        """根据路径获取到所有的 image 文件名

        :param root_dir: data路径
        :param label_dir: label
        """
        self.root_dir = root_dir
        self.label_dir = label_dir
        self.data_dir = os.path.join(self.root_dir, self.label_dir)
        self.img_names = os.listdir(self.data_dir)

    def __getitem__(self, index) -> T_co:
        """overwrite 后才可以通过下标访问 dataset

        :param index: 下标
        :return: img(PIL), label
        """
        img_name = self.img_names[index]
        img_path = os.path.join(self.data_dir, img_name)
        img = Image.open(img_path)
        label = self.label_dir
        return img, label  # 放回什么数据由自己定义

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

使用MyDataset类

在交互模式中导入MyDataset并执行如下语句：

ants_dataset = MyDataset("dataset/train", "ants")  # 创建dataset
bees_dataset = MyDataset("dataset/train", "bees")

img, label = ants_dataset[47]  # 通过下标访问数据
len(bees_dataset)  # 获取dataset长度
img.show()  # 显示图片

dataset = ants_dataset + bees_dataset  # 将dataset相加

Tensorboard

使用tensorboard可以按照日志的形式记录训练模型时产生的一些数据（标量、图片等），同时也可以对记录的数据进行可视化，方便我们对现有模型进行分析。例如，借助tensorboard记录loss 函数随着训练训练周期的下降过程并可视化。

使用Tensorboard首先要在conda环境中通过如下命令进行安装：

conda install tensorboard

记录标量（scalar）

记录标量需要用到torch.utils.tensorboard中的SummaryWriter类的add_scalar方法。

创建如下python脚本并执行：

from torch.utils.tensorboard import SummaryWriter

# 参数：log_dir，日志存放的路径名称
writer = SummaryWriter("logs")

# 记录曲线 y = x^2
# 参数一：tag，数据标识符
# 参数二：scalar_value，标量值，y轴数据
# 参数三：global_step，步骤，x轴数据
for i in range(100):
    writer.add_scalar("y = x^2", i ** 2, i)
# 调用close() 确保数据刷入磁盘
writer.close()

脚本执行成功后会在脚本文件所在当前路径下创建logs文件夹并生成日志文件。日志生成成功后在conda环境中执行以下命令来启动tensorboard可视化服务：

# 指明log存放的位置以及服务启动时使用的端口，默认端口为6006
# 多人使用同一台服务器时应该避免以下两个参数值相同
tensorboard --logdir=second/logs --port=6008

该命令执行成功后会给出如下图中的提示：

在浏览器中打开以上链接即可查看可视化情况：

注意：Tensorboard是根据不同的tag来进行可视化的，如果多次写入的日志有相同tag，可能会导致可视化时出问题。只需删除日志并重新生成，然后重启可视化服务即可解决。

记录图像（image）

记录图像可以使用torch.utils.tensorboard中的SummaryWriter类的add_image方法。

该方法定义如下：

def add_image(self, tag, img_tensor, global_step=None, walltime=None, dataformats='CHW')
# tag(string) 同上
# img_tensor (torch.Tensor, numpy.array, or string/blobname): 图像数据
# global_step (int): 步骤
# dataformats (string)：图像数据格式，不为CHW则需要指明

创建如下python脚本并执行：

from torch.utils.tensorboard import SummaryWriter
from PIL import Image
import numpy as np

writer = SummaryWriter("logs")
img_path = "E:\\code\\python\\pytorch-learn\\dataset\\train\\ants\\0013035.jpg"
# img 为 PIL.JpegImagePlugin.JpegImageFile 类型
img = Image.open(img_path)
img_array = np.array(img)

writer.add_image("ants", img_array, 0, dataformats="HWC")
writer.close()

脚本执行成功后按之前同样的方式启动可视化服务即可。

如果日志中同一个tag标识的数据有多个global_step则可以如下图所示：拖动轴来查看global_step之间的图像变化。

Transforms

transforms是位于torchvision包中的一个模块，模块中有多个类可以对将要输入神经网络的数据进行转换以适应网络的需要。

ToTensor

该类可以将 PIL Image or numpy.ndarray 转换为张量，以方便图像输入到神经网络中，如下代码做了简单示例：将PIL Image 转换为张量后借助tensorboard直接写入日志文件：

from torchvision import transforms
from PIL import Image
from torch.utils.tensorboard import SummaryWriter
# 创建 transforms 模块中的 ToTensor 类
trans_to_tensor = transforms.ToTensor()

# 获取 PIL 图像
img_path = "E:/code/python/pytorch-learn/dataset/train/ants" \
           "/6743948_2b8c096dda.jpg "
img = Image.open(img_path)

# 将 PIL 图像转化为 tensor 类型
#   1. 将输入的数据shape W，H，C ——> C，W，H
#   2. 将所有数除以255，将数据归一化到[0，1]
# 实例像方法一样调用，这里实际调用的是ToTensor类中的__call__方法，
#   更多：https://zhuanlan.zhihu.com/p/370234492
img_tensor = trans_to_tensor(img)

# 使用 tensorboard 将 tensor 图像写入日志
writer = SummaryWriter("logs")
writer.add_image("img_tensor", img_tensor, 1)
writer.close()

启动tensorbard可视化服务并在浏览器中访问：

Normalize

normalize类可以利用均值和标准差对tensor图像进行归一化。该类构造函数定义如下：

# mean (sequence): 代表图像每个通道均值的序列
# std (sequence): 代表图像每个通道标准差的序列
# inplace(bool,optional): 是否原地修改tensor
def __init__(self, mean, std, inplace=False):

该类的具体归一化操作如下：

output[channel] = (input[channel] - mean[channel]) / std[channel]

拿任意通道来说，如果原值 value 为 [0, 1]，传入的 mean 为 0.5，std 为 0.5 那么归一化操作即为：(value - 0.5) / 0.5 = 2*value - 1，带入value范围后计算得到value的新范围为[-1, 1]。

创建如下脚本并执行：

from torchvision import transforms
from PIL import Image
from torch.utils.tensorboard import SummaryWriter
# 创建 transforms 模块中的 ToTensor 类
trans_to_tensor = transforms.ToTensor()
# 创建 transforms 模块中的 Normalize 类，三个通道的mean和std都为0.5
trans_normalize = transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])

# 获取 PIL 图像
img_path = "E:/code/python/pytorch-learn/dataset/train/ants" \
           "/6743948_2b8c096dda.jpg "
img = Image.open(img_path)

# 将 PIL 图像转化为 tensor 类型
#   1. 将输入的数据shape W，H，C ——> C，W，H
#   2. 将所有数除以255，将数据归一化到[0，1]
img_tensor = trans_to_tensor(img)
# 将 tensor 进行归一化
img_normalize = trans_normalize(img_tensor)
print(type(img_normalize))
# <class 'torch.tensor'="">

# 使用 tensorboard 将 tensor 图像写入日志
writer = SummaryWriter("logs")
writer.add_image("img_normalize", img_normalize, 1)
writer.close()

启动tensorbard可视化服务并在浏览器中访问：

Resize

Resize类可对tensor或PIL图像进行缩放，构造函数定义如下：

# size (sequence or int): 所需的输出大小
#   如果size是形如（h, w）的高和宽的序列，输出时就按此匹配；
#   如果size是一个int，则选择小的边来进行等比例缩放，如果 height > width，
#   则缩放后为(size * height / width, size)
def __init__(self, size, interpolation=InterpolationMode.BILINEAR, max_size=None, antialias=None):

创建如下脚本并执行：

from torchvision import transforms
from PIL import Image
from torch.utils.tensorboard import SummaryWriter
# 创建到 transforms 模块中的 ToTensor 类
trans_to_tensor = transforms.ToTensor()
# 创建 transforms 模块中的 Normalize 类
trans_normalize = transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
# 创建 transforms 模块中的 Resize 类
trans_resize_52x52 = transforms.Resize((80, 40))
trans_resize_52 = transforms.Resize(200)

# 获取 PIL 图像
img_path = "E:/code/python/pytorch-learn/dataset/train/ants" \
           "/6743948_2b8c096dda.jpg "
img = Image.open(img_path)

# 使用 tensorboard 记录多个处理步骤的图像
writer = SummaryWriter("logs")

# 将 PIL 图像转化为 tensor 类型
#   1. 将输入的数据shape W，H，C ——> C，W，H
#   2. 将所有数除以255，将数据归一化到[0，1]
img_tensor = trans_to_tensor(img)
# 步骤0：记录tensor图像
writer.add_image("img-tensor-normalize-resize_52x52_resize_52", img_tensor, 0)

# 将 tensor 进行归一化
img_normalize = trans_normalize(img_tensor)
# 步骤1：记录normalize后的图像
writer.add_image("img-tensor-normalize-resize_52x52_resize_52",
                 img_normalize, 1)

# 对normalize后的tensor进行(52, 52)的resize
img_resize_52x52 = trans_resize_52x52(img_normalize)
# 步骤2：记录(52, 52)resize后的图像
writer.add_image("img-tensor-normalize-resize_52x52_resize_52",
                 img_resize_52x52, 2)

# 对(52, 52)resize后的tensor再次进行(52)的resize
img_resize_52 = trans_resize_52(img_resize_52x52)
# 步骤3：记录(52)resize后的图像
writer.add_image("img-tensor-normalize-resize_52x52_resize_52",
                 img_resize_52, 3)
writer.close()

启动tensorbard可视化服务并在浏览器中访问：

Compose

在上一节的示例中，我们对图像的处理是基于一系列transform，为了记录每个transform执行后的图像，我们在每两个transform之间插入了记录图像操作。如果我们不需要记录每个transform执行后的图像，那么我们可以通过Compose来顺序执行一系列transform。

Compose类的构造函数定义如下：

# transforms (list of Transform objects):
#   需要顺序执行的transfrom对象组成的list
def __init__(self, transforms):

该类的核心思想位于__call__函数中：

def __call__(self, img):
    for t in self.transforms:
        img = t(img)
    return img

有了compose后，为了得到上一节最后的结果，我们可以对上一节的脚本做如下简化：

from torchvision import transforms
from PIL import Image
from torch.utils.tensorboard import SummaryWriter
# 创建到 transforms 模块中的 ToTensor 类
trans_to_tensor = transforms.ToTensor()
# 创建 transforms 模块中的 Normalize 类
trans_normalize = transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
# 创建 transforms 模块中的 Resize 类
trans_resize_80x40 = transforms.Resize((80, 40))
trans_resize_200 = transforms.Resize(200)
# 创建 transforms 模块中的 Compose 类
trans_compose = transforms.Compose([trans_to_tensor, trans_normalize,
                                    trans_resize_80x40, trans_resize_200])

# 获取 PIL 图像
img_path = "E:/code/python/pytorch-learn/dataset/train/ants" \
           "/6743948_2b8c096dda.jpg "
img = Image.open(img_path)

# 执行 Compose
img_compose = trans_compose(img)

# 记录compose后的图像
writer = SummaryWriter("logs")
writer.add_image("img-tensor-normalize-resize_80x40-resize_200-compose",
                 img_compose, 0)
writer.close()

启动tensorbard可视化服务并在浏览器中访问：

transforms模块中还有很多其他的transform类，这些类在需要时查询文档即可。查询文档时，重点要弄清楚类或函数的作用以及输入输出类型。

Dataset 和 Transforms结合使用

在pytorch官网首页，按不同模块分别提供了文档，包括：核心模块，音频模块，文本模块，视觉模块等：

在torchvision.datasets中，提供了对常用数据集的支持，它可以帮助我们下载、解压并快速使用数据集：

以CIFAR10数据集为例，根据如下文档使用即可：

根据文档可以知道，要使用CIFAR10数据集则使用torchvision.datasets.CIFAR10类，该类的构造函数要求传入如下参数：

# root (string)：数据集所在目录
# train (bool, optional)：为True，创建训练集；为false，创建测试集
# transform (callable, optional)：处理 PIL image 的function/transform
# target_transform (callable, optional)：处理 target（图像类别）的function/transform
# download (bool, optional)：为true则下载数据集到root目录中，如果已经存在则不会下载
def __init__(
    self,
    root: str,
    train: bool = True,
    transform: Optional[Callable] = None,
    target_transform: Optional[Callable] = None,
    download: bool = False,
) -> None:

根据文档中的__getitem__函数我们知道，通过下标访问Dataset时会返回(image, target)，即图像和图像的种类。

创建如下脚本并执行（使用训练集）：

import torchvision

# 创建transform，将 PIL Image 转换为 tensor
data_trans = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor()
])

# 建议download始终为True，即使你自己提前下载好了数据集
# 下载慢的话可以拷贝控制台输出的下载地址，然后到迅雷下载好后再将压缩包拷贝至root下即可
# 下载地址也可以在torchvision.datasets.CIFAR10类的源码中查看
train_set = torchvision.datasets.CIFAR10(root="./dataset", train=True,
                                         transform=data_trans, download=True)

img, target = train_set[0]
print(type(img))  # <class 'torch.tensor'="">
print(target)     # 6
# ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
print(train_set.classes)
print(train_set.classes[target])  # frog

创建如下脚本并执行（使用测试集）：

import torchvision

test_set = torchvision.datasets.CIFAR10(root="./dataset", train=False,
                                        download=True)

img, target = test_set[0]
img.show()
print(type(img))  # <class 'pil.image.image'="">
print(target)     # 3
print(test_set.classes[target])  # cat

DataLoader

借用DataLoader可以将数据输入神经网络，查看官方文档：

根据官方文档我们知道，在创建DataLoader时需要的主要的参数如下：

# dataset (Dataset): Dataset类型数据集
# batch_size (int, optional): 批大小，default: 1
# shuffle (bool, optional): 每个 epoch 后是否打乱数据，default: False
# num_workers (int, optional): 加载数据时使用的子进程数量，为0表示使用主进程加载，default: 0。可设置为你的cpu核心数
# drop_last (bool, optional): 当数据个数不能被batch_size整除时，为True表示丢弃最后一批，为False表示不丢弃，default: False
def __init__(self, dataset: Dataset[T_co], batch_size: Optional[int] = 1,
                 shuffle: bool = False, sampler: Optional[Sampler] = None,
                 batch_sampler: Optional[Sampler[Sequence]] = None,
                 num_workers: int = 0, collate_fn: Optional[_collate_fn_t] = None,
                 pin_memory: bool = False, drop_last: bool = False,
                 timeout: float = 0, worker_init_fn: Optional[_worker_init_fn_t] = None,
                 multiprocessing_context=None, generator=None,
                 *, prefetch_factor: int = 2,
                 persistent_workers: bool = False):

创建如下脚本并执行：

import torchvision
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

def train():
    trans_to_tensor = torchvision.transforms.ToTensor()
    train_dataset = torchvision.datasets.CIFAR10(root="./dataset",
                                                 transform=trans_to_tensor,
                                                 download=True)
    # dataloader 按照batch_size打包img和target，如图所示
    train_dataloader = DataLoader(dataset=train_dataset, batch_size=64,
                                  shuffle=False, num_workers=16,
                                  drop_last=False)

    writer = SummaryWriter("logs")

    for epoch in range(2):
        for batch_num, batch_data in enumerate(train_dataloader):
            # print(epoch, " : ", batch_num)
            batch_img, batch_target = batch_data
            writer.add_images(
                "CIFAR10-batch64-no_shuffle-no_drop_last_epoch{}".format(epoch)
                , batch_img, batch_num)

    writer.close()

if __name__ == "__main__":
    train()

在windows中如果num_workers > 1报错，解决方法如下：

1.修改代码为如下格式：

def train():
    # Here was inserted the whole code that train the network ...
if __name__ == '__main__':
    train()

2.如果代码修改后仍然报错：OSError: [WinError 1455] 页面文件太小,无法完成操作。则还需要调整你的python环境所在盘的虚拟内存大小:

进入高级系统设置：

点击：高级 -> 性能 -> 设置：

点击：高级 -> 更改：

按照下图，将D盘的虚拟内存调整为系统管理的大小（重启后生效）。任然报错就自定义大小为一个很大的值，如：10GB-100GB。

脚本执行成功后，启动tensorbard可视化服务并在浏览器中访问：

step=170时：（日志比较大，浏览器中可能需要多次刷新才能完全加载）

step=781时：

由于我们在创建DataLoader时，参数shuffle=False使得epoch0和epoch1之间相同批次抽取的图像是相同的；参数drop_last=False使得即使最后一个批次图像不足64也没有被舍弃。

修改脚本使得：shuffle=True，drop_last=True，然后再次执行脚本并启动tensorbard可视化服务。浏览器中访问结果如下：

如图所示，epoch0和epoch1之间相同批次（780）抽取的图像不相同，且第781个批次被舍弃。