第四章 PyTorch 数据处理工具箱

4.1 数据处理工具箱概述

"""
<===== 4.1 数据处理工具箱概述 =====>
torch.utils.data 是 PyTorch 中用于数据处理的核心工具箱，提供了高效的数据加载和预处理功能。它包括以下 4 个类:
1. Dataset: 用于表示数据集的抽象类，其他数据集需要继承这个类。用户可以继承该类并实现 __len__ 和 __getitem__ 方法来定义自己的数据集。
2. DataLoader: 定义一个新的迭代器，用于将 Dataset 封装成一个可迭代的数据加载器，支持批量加载数据(Batch)、打乱数据顺序(Shuffle)、多线程加载等功能。
3. random_split: 用于将数据集随机划分为给定长度非重叠的多个子集，常用于训练集和验证集的划分。
4. Sampler: 定义了从数据集中采样样本的策略，可以自定义采样方式，如随机采样、顺序采样等。

torchvision 是 PyTorch 中的计算机视觉工具箱，提供了常用的数据集、模型和图像处理工具。它包括以下 3 个模块:
1. datasets: 提供了常用的计算机视觉数据集，如 MNIST、CIFAR-10、ImageNet 等，方便用户直接加载和使用。设计上是继承自 torch.utils.data.Dataset 类。
2. models: 提供了预训练的深度学习模型，如 ResNet、VGG、AlexNet 等，用户可以直接使用这些模型进行迁移学习或微调。
3. transforms: 提供了常用的图像预处理和数据增强方法，如裁剪、缩放、归一化等，方便用户对图像数据进行处理。主要包括对 PIL 图像和 Tensor 的转换操作。
4. utils: 包含两个函数，make_grid 用于将多张图像拼接成一张网格图像，save_image 用于将 Tensor 保存为图像。
"""

4.2 utils.data

import torch
import torch.utils.data as Data
import numpy as np

# <====== 4.2 utils.data ======>
"""
utils.data 包括 Dataset 和 DataLoader。
torch.utils.data.Dataset 是一个抽象类，表示数据集。所有自定义数据集都应该继承这个类，并重写以下两个方法：__len__ 和 __getitem__。
__len__ 方法返回数据集的大小，__getitem__ 方法根据索引返回数据集中的一个样本。
__getitem__ 方法可以返回任何类型的数据，但通常返回一个 (input, target) 元组，分别表示输入数据和对应的标签。
所以通过 torch.utils.data.Dataloader 来定义一个新的迭代器，实现批量读取。
"""


# 定义获取数据集的类: 该类继承基类 Dataset，自定义一个数据集及其对应标签。
class TestDataset(Data.Dataset):  # 继承 Dataset
    def __init__(self):
        self.Data = np.array([[1, 2], [3, 4], [2, 1], [3, 4], [4, 5]])  # 一些由 2 维向量表示的数据集
        self.Label = np.array([0, 1, 0, 1, 2])  # 每个数据集对应的标签

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

    def __getitem__(self, index):  # 传入索引 index，返回数据和标签
        # 把 Numpy 转换为 Tensor
        txt = torch.from_numpy(self.Data[index])  # 取出数据集
        label = torch.tensor(self.Label[index])  # 取出对应标签
        return txt, label  # 返回数据集和标签


# 获取数据集中的数据
Test = TestDataset()
print(f'获取第三个元素: {Test[2]}')  # 相当于调用 __getitem__(2)
print(f'数据集的长度: {len(Test)}')  # 相当于 Test.__len__()

"""
以上的数据由元组格式返回，每次只返回一个样本。如果希望批量处理，同时进行数据打乱和并行加速等操作，可以选择 DataLoader。
utils.data.DataLoader(
    dataset,  # 加载的数据集
    batch_size=1,  # 每个 batch 包含多少个样本
    shuffle=False,  # 是否打乱数据
    sampler=None,  # 定义从数据集中提取样本的策略
    batch_sampler=None,  # 类似于 sampler，但返回的是一批样本的索引
    num_workers=0,  # 使用多少个子进程来加载数据，0 表示不使用多进程
    collate_fn=None,  # 如何将多个样本合并成一个 batch
    pin_memory=False,  # 是否将数据放到 CUDA 的固定内存中
    drop_last=False,  # 如果数据集大小不能被 batch_size 整除，是否丢弃最后一个不完整的 batch
    timeout=0,  # 如果数据加载时间超过 timeout 秒，则抛出异常
    worker_init_fn=None,  # 每个子进程初始化时调用的函数
)
"""

# 定义一个 DataLoader
test_loader = Data.DataLoader(dataset=Test, batch_size=2, shuffle=False, num_workers=0)
for i, train_data in enumerate(test_loader):
    print(f'第 {i} 个 batch:\n{train_data}')  # train_data 是一个 batch 的数据
    data, label = train_data
    print(f'数据: {data}')  # 每个 batch 的数据
    print(f'标签: {label}')  # 每个 batch 的标签

# 一般用 data.Dataset 处理同一个目录下的数据。如果数据在不同目录下（不同目录代表不同类别），Dataset 就不是非常方便。可以使用 torchvision 包来处理。

4.3 torchvision

import torch
import numpy as np
import torch.utils.data as Data
from torchvision import transforms, datasets, utils
from PIL import Image
import matplotlib.pyplot as plt

# <====== 4.3 torchvision ======>
"""
torchvision 有 4 个子模块：datasets、models、transforms 和 utils。model 将在后面讲解，在 3.4 中已经使用过 datasets 下载过数据集。
现在主要介绍如何使用 transforms 对原数据进行预处理、增强等，以及 datasets 下的 ImageFolder 处理自定义数据。
"""

# <------ 4.3.1 transforms ------>
"""
transforms 提供了对 PIL Image 和 Tensor 对象的常用操作。
A 对 PIL Image 的操作:
1) Resize(size, interpolation=2): 调整图像大小，size 可以是一个整数或一个 (width, height) 元组。
2) CenterCrop(size): 从图像中心裁剪出指定大小的区域，size 可以是一个整数或一个 (width, height) 元组。
RandomCrop 随机裁剪，RandomResizedCrop 裁剪随机大小。
3) Pad(padding, fill=0, padding_mode='constant'): 对图像进行填充，padding 可以是一个整数或一个 (left, top, right, bottom) 元组。
4) ToTensor(): 将 PIL Image 或 numpy.ndarray 转换为 Tensor，并将像素值归一化到 [0, 1] 范围内。如把 (H x W x C) 的图像转换为 (C x H x W) 的 Tensor。
5) RandomHorizontalFlip(p=0.5): 以概率 p 随机水平翻转图像。
6) RandomVerticalFlip(p=0.5): 以概率 p 随机垂直翻转图像。
7) ColorJitter(brightness=0, contrast=0, saturation=0, hue=0): 随机改变图像的亮度、对比度、饱和度和色调。
B 对 Tensor 的操作:
1) Normalize(mean, std, inplace=False): 使用均值和标准差对 Tensor 进行归一化。mean 和 std 应该是与图像通道数相同的列表。
2) ToPILImage(mode=None): 将 Tensor 或 numpy.ndarray 转换为 PIL Image。

如果要对数据集进行多个变换，可以使用 transforms.Compose() 将多个变换组合在一起，按顺序依次执行。类似于 nn.Sequential()。
"""
transforms.Compose([
    # 将给定的 PIK.Image 进行中心切割，得到给定的 size,
    # 可以是元组 (target_height, target_width) 也可以是一个整数，这种情况下切割出来一个正方形
    transforms.CenterCrop(10),  # 10 x 10
    transforms.RandomCrop(20, padding=0),  # 随机选取切割点中心位置，20 x 20, 并在每个边上填充 0 个像素
    transforms.ToTensor(),  # 把一个取值为[0, 255]的PIL.Image或者形状为(H,W,C)的np.ndarry 转换为取值为[0.0, 1.0]，形状为 (C,H,W)的torch.FloatTensor
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])  # 对每个通道进行归一化到[-1, 1]，均值为0.5，标准差为0.5
])
# 还可以自己定义一个 Lambda 函数进行变换
transforms.Lambda(lambda x: x + 10)  # 对输入的每个像素值加 10


# <------ 4.3.2 ImageFolder ------>
"""
当文件依据标签处于不同文件夹下时，如:
————— data
     |———— label1
     |    |————— 001.jpg
     |    |————— 002.jpg
     |———— label2
     |    |————— 001.jpg
     |    |————— 002.jpg
     .......
就可以利用 torchvision.datasets.ImageFolder 来直接构造出 dataset，代码如下:
dataset = datasets.ImageFolder(Path)
loader = utils.data.DataLoader(dataset)
ImageFolder 会自动为每个子文件夹分配一个标签，按字母顺序排序，从 0 开始编号。这样载入 DataLoader 后，每个 batch 的数据就会包含图像和对应的标签。
"""
# 利用 ImageFolder 读取不同目录下的数据，然后用 transforms 进行图像预处理，然后使用 DataLoader 加载数据，将处理后的数据用 torchvision.utils
# 中的 save_image 保存为 png 文件，然后用 Image.open 打开这个图片。
# 我随意找了 5 张猫和狗的图片，放在 data/cat 和 data/dog 目录下。
my_transforms = transforms.Compose([
    transforms.RandomResizedCrop(224),  # 随机裁剪为 224 x 224
    transforms.RandomHorizontalFlip(),  # 随机水平翻转
    transforms.ToTensor(),  # 转换为 Tensor
])
train_data = datasets.ImageFolder(root='./data', transform=my_transforms)  # 读取数据并预处理
train_loader = Data.DataLoader(dataset=train_data, batch_size=6, shuffle=True)  # 加载数据

for i_batch, img in enumerate(train_loader):
    if i_batch == 0:
        print(img[1])  # e.g.tensor([1, 0, 1, 1, 0, 0])。img[1] 是标签, img[0] 是图像数据
        fig = plt.figure()
        grid = utils.make_grid(img[0])  # 将一个 batch 的图像拼接成一张网格图像
        plt.imshow(grid.numpy().transpose((1, 2, 0)))  # 转换为 (H, W, C) 格式
        plt.show()
        utils.save_image(grid, 'cat_dog.png')  # 保存图像
        break

img = Image.open('cat_dog.png')  # 打开保存的图像