前言

本文的定位如同本站的另外一篇博客机器学习通用技术：sklearn快速入门，在学习机器学习与深度学习的过程中，会涉及到实现这些算法的编程工具，在机器学习中是sklearn，在深度学习中，目前的主流工具是Pytorch。这些工具的学习与算法的学习相辅相成，但是工具中有一部分常用的通用性内容，故单开一文，可作为Pytorch的快速入门，迅速跑起来自己的第一个模型，也可以作为参考资料，以便遗忘时翻阅。

Pytorch构建模型

首先，我们需要学习如何在Pytorch中构建一个神经网络模型，Pytorch中关于构建神经网络的包为torch.nn，import时一般写为

1
from torch import nn

相比于实现机器学习的skalern库，Pytorch构建神经网络模型的过程会更复杂一点，这主要是因为，神经网络通常是由多个不同的网络层组成的，可以灵活的调整网络层的顺序，参数以及组合，就像搭积木一样，因此，torch.nn中实现了大多数常用的网络层，并且可以通过特定的方式将其进行组合，下面介绍两种主要的构建模型的方法

nn.Sequential

nn.Sequential()是torch.nn中提供的一个顺序容器，其用法也非常简单，将网络层按顺序填入该容器中即可构建神经网络模型，例如

1
2
3
4
model = nn.Sequential(
    nn.Flatten(),  # 展平层
    nn.Linear(784, 10)  # 线性层
)

这是一种非常直观的构建模型方法，适用于构建简单的神经网络。但是其缺点也在于此，由于其只能顺序堆叠，对于一些特殊的网络设计，例如ResNet中的残差设计，就不能使用该方法进行构建。

model类

通过model类来构建神经网络模型是Pytorch中最常用，也是最通用的方法。顾名思义，这是一个类，这个类需要继承nn.Module，并且实现__init__() 和 forward() 两个方法，例如

1
2
3
4
5
6
7
8
9
10
class LR(nn.Module):
    def __init__(self):
        super().__init__()
        self.input = nn.Flatten()  # 展平层
        self.L = nn.Linear(784, 10)  # 线性层
        
    def forward(self, x):
        x = self.input(x)
        x = self.L(x)
        return x

首先，在__init__() 方法中，必须调用父类初始化，即super().__init__()，一般也会将各网络层的定义写在该方法中，不过这只是习惯，并不是必须的。而 forward() 表示前向传播，其接受神经网络的输入，并且返回神经网络的输出，例如上述代码中，输入x经过展平层与线性层，得到神经网络的输出并返回。当然了，由于这是一个类，所以还需要对这个类进行实例化才算完成模型的构建，即

1
model = LR()

注：初学者可能会对这种构建方法一头雾水，但这是不得不学习的，我能分享的经验就是，最开始只需要依葫芦画瓢，写多了也就逐渐理解了，事物都是螺旋上升的嘛。

Pytorch训练模型

然后，是如何使用Pytorch训练模型，训练一个神经网络需要准备模型，数据，损失函数以及优化算法，模型的构建我们在上一小节就完成了，而数据的读取则是另一个复杂的话题，我们会在后文进行讲解，本小节仅介绍如何使用Pytorch中自带的数据集模块

注：在Pytorch自带的数据集模块中，集成了许多深度学习的经典数据集，对于新手学习算法来说，使用这些数据集就足够了。因此，关于读取数据集的通用方法，本文选择放在后面的小节进行讲解

torchvision.datasets

torchvision.datasets是Pytorch中自带的数据集模块，其中包含了大量有关计算机视觉的经典数据集，例如MNIST，CIFAR10等，关于该模块包含了哪些数据集，可以参考官方文档或者官方文档中文翻译版，接下来我们将介绍如何使用该模块读取这些经典的数据集

注1：从模块名称不难看出，该数据集模块其实是归属于torchvision库的，这是一个基于Pytorch开发的，有关计算机视觉的库，其提供了许多使用Pytroch做有关视觉的深度学习时常用的工具，例如经典的视觉数据集，用于跑模型demo，以及经典的有关计算机视觉的深度学习模型的实现，例如ResNet系列，一行代码就可以调用这些模型，甚至可以调用使用ImageNet预训练好的模型进行迁移学习，并且其源码都是用python写的，也是非常好的学习材料。但是，在安装Pytorch时，默认会顺便安装torchvision，这点从Pytroch官网提供的安装命令也能看得出来，因此，将其称为Pytorch自带的数据集模块也是蛮合理的吧。

注2：类似于torchvision这样的库还有不少，其都是Pytorch生态的组成部分，可以在PyTorch官网上查看，例如用于NLP领域的torchtext（可惜已经停止开发了），以及用于强化学习的torchrl等等。

首先，我们需要import该模块，一般只import所需的数据集，以MNIST为例

1
from torchvision.datasets import MNIST

除此之外我们还需要import两个模块

1
2
from torch.utils.data import DataLoader  # 数据读取器
from torchvision import transforms  # 数据预处理

其中，DataLoader是torch自带的数据读取器，用于将数据集（严格来说是Dataset类）转换为一个python迭代器，以便于在训练过程中进行小批量读取，而transforms是torchvision中提供的用于处理图像数据的模块，其功能非常强大，不过在这里其作用仅仅是将图像数据转换为torch的tensor类型，完整的数据集读取代码如下

1
2
3
4
5
6
# 读取数据集
data_MNIST_train = MNIST('Data/mnist_data', train=True, download=True, transform=transforms.ToTensor())
data_MNIST_test = MNIST('Data/mnist_data', train=False, download=True, transform=transforms.ToTensor())
# 生成数据读取器
MNIST_train = DataLoader(data_MNIST_train, batch_size=64, shuffle=True)  # shuffle表示打乱
MNIST_test = DataLoader(data_MNIST_test, batch_size=64, shuffle=True)

其中，数据集函数MNIST的参数含义从左至右分别为，数据集存储路径，是否为训练集（反之为测试集），是否需要下载（第一次调用需要下载），对数据集所做的预处理（此处仅需要将数据类型转化为tensor即可）。而DataLoader的参数含义从左至右分别为，数据集（Dataset类），批量大小，是否打乱，至此，我们就得到了两个python迭代器。

注：对于python生成器不熟悉的读者可以参考Python：迭代器与生成器，其中介绍了python中的迭代器与生成器，以及如何使用python生成器进行小批量读取

损失函数与优化算法

Pytorch中，损失函数位于torch.nn中，而优化算法位于torch.optim中，两者均以类的形式存在，在使用时进行实例化即可，且损失函数与优化算法的选择也相对比较固定，一般来说，对于分类问题选择交叉熵损失函数，对于回归问题选择MSE损失函数，而优化算法一般在SGD与Adam中进行选择，此处我们选择使用交叉熵与SGD，代码如下

1
2
3
4
import torch
# 损失函数与优化算法
criterion = torch.nn.CrossEntropyLoss()  # 交叉熵
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)  # SGD

上述代码中，SGD函数的参数model.parameters()代表模型的可学习参数，lr代表学习率。

注：对于梯度下降不熟悉的读者可以参考深度学习基础：梯度下降

训练流程与GPU加速

准备工作完成，现在可以正式开始模型的训练了，不过此处需要补充一个非常重要的内容，就是如何调用GPU进行加速，此处仅考虑Windows系统下使用单张NVIDIA的显卡进行模型训练，并且已经安装了正确的Pytorch及CUDA环境的情况，使用下面的代码可以输出可使用的显卡设备，如果没有可以使用的显卡设备，则输出为CPU

1
2
# 训练设备
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  # 调用GPU进行加速，device为GPU设备名，可以使用.to(device)将模型与数据放到GPU上

注1：虽然在深度学习中GPU加速很重要，但是至少本文的内容用CPU还是跑得起来的，而且本文的代码也兼容了使用CPU运行的情况，不过真的要学习深度学习的话还是建议搞张N卡来跑模型，或者用Apple的M系列芯片来跑也是不错的选择（不过M系列芯片得改一下代码

注2：本文并不是一篇关于环境配置的文章，配深度学习的环境还是有点小复杂的，如果环境出了什么问题，没办法顺利调用GPU的话，请出门左转，配置好环境再继续阅读，或者就将就着用CPU先

然后可以使用to()方法，将模型或者数据放到指定的设备上，这里先将模型放到GPU上

1
model.to(device)  # 将模型放到device上

现在可以正式进入模型训练了，下面我们给出模型训练的核心代码

1
2
3
4
5
6
7
8
9
10
11
model.train()  # 调整为训练模式

for data, target in MNIST_train:
    data, target = data.to(device), target.to(device)  # 将数据放到GPU上
        
    optimizer.zero_grad()  # 梯度清零
        
    output = model(data)  # 前向传播
    loss = criterion(output, target)  # 计算损失值
    loss.backward()  # 反向传播
    optimizer.step()  # 更新参数

首先，我们使用train()方法令模型进入训练模式，该模式的主要作用是让如dropout一类，在训练过程中开启，在推理过程中关闭的网络层可以正确的工作，在此处其实并无作用，不过为了代码的通用性，还是加上。然后是使用for循环配合生成器，对数据进行小批量读取，并且将读取的数据放到GPU上。接着对梯度进行清零，因为Pytorch中，优化器的梯度是会累计的，因此在每一轮训练之前都需要手动清零。再之后便是进行前向传播得到模型输出，然后计算损失值，进行反向传播，最后使用优化算法更新参数，每一步对应一行代码，大多数情况下这几行代码都是不会变的。

注：模型和数据必须放在同一个训练设备上，否则会报错

最后，将上述代码封装为函数，并简单加上一些信息打印，就实现了一个可以进行一个epoch训练的函数，代码如下

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
def train(epoch):
    model.train()  # 调整为训练模式
    run_loss = 0.0
    
    for data, target in MNIST_train:
        data, target = data.to(device), target.to(device)  # 将数据放到GPU上
        
        optimizer.zero_grad()  # 梯度清零
        
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        
        run_loss += loss.item()  # 存储本mini_batch的损失值
        
    train_loss = run_loss / len(MNIST_train)  # 将损失值之和除以样本量，用于近似本epoch的损失值
    print('Epoch: {} \tTraining Loss: {:.6f}'.format(epoch, train_loss))

调用上述函数进行10个epoch的训练

1
2
for epoch in range(10):
    train(epoch)
1
2
3
4
5
6
7
8
9
10
Epoch: 0 	Training Loss: 0.479043
Epoch: 1 	Training Loss: 0.337416
Epoch: 2 	Training Loss: 0.314644
Epoch: 3 	Training Loss: 0.302373
Epoch: 4 	Training Loss: 0.294522
Epoch: 5 	Training Loss: 0.288994
Epoch: 6 	Training Loss: 0.284427
Epoch: 7 	Training Loss: 0.281150
Epoch: 8 	Training Loss: 0.278040
Epoch: 9 	Training Loss: 0.275650

至此，我们就使用Pytorch完成了第一个神经网络模型的训练。

模型测试

接着就是模型的测试了，这个就很简单了，只需要计算模型输出，再计算准确率即可，不过分类模型的输出是概率向量，因此需要将其转换为标签，另外，模型测试时需要使用eval()方法将模型调整为测试模型，该模式与上面的训练模式对应，最后，跑测试的时候可以关闭梯度计算，提高计算速度，具体代码如下

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
def test():       
    model.eval()  # 调整为测试模式
    correct = 0
    total = 0
    
    with torch.no_grad():  # 不需要计算梯度
        for data, target in MNIST_test:
            data, target = data.to(device), target.to(device)  # 将数据放到GPU上
            
            output = model(data)  # 计算输出
            
            predicted = torch.argmax(output, 1)  # 向量转换为标签
            
            total += target.size(0)  # 测试集数量
            correct += (predicted == target).sum().item()  # 预测值与正确类别是否相等，相等则+1（labels是类别数）
            
    print('Accuracy on test set: %d %%' % (100 * correct / total))  # 正确率

调用上述函数进行模型测试

1
test()  # 输出为: Accuracy on test set: 92 %

注：经过上面的学习，不难发现，Pytorch训练模型是一个相对固定的流程，因此大多数时候大家都会根据自己的使用习惯写一个训练函数，包含训练，测试，结果打印，画图等功能，然后就一直用下去

模型保存与加载

最后，我们来介绍模型的保存，Pytorch的模型保存非常的方便，使用torch.save函数即可实现，常用的有两种保存方式，一种是保存整个模型，另一种是只保存模型的参数，我们一个一个来看。首先是保存整个模型

1
2
# 保存整个模型
torch.save(model, 'model.pt')

使用torch.save函数直接保存整个模型，保存的模型文件后缀通常为.pt或者.pth。而模型的读取也非常简单，使用torch.load函数即可直接读取

1
2
# 加载整个模型
loaded_model = torch.load('model.pt')
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
loaded_model.eval()  # 调整为测试模式
correct = 0
total = 0
    
with torch.no_grad():  # 不需要计算梯度
    for data, target in MNIST_test:
        data, target = data.to(device), target.to(device)  # 将数据放到GPU上
            
        output = model(data)  # 计算输出
            
        predicted = torch.argmax(output, 1)  # 向量转换为标签
            
        total += target.size(0)  # 测试集数量
        correct += (predicted == target).sum().item()  # 预测值与正确类别是否相等，相等则+1（labels是类别数）
            
print('Accuracy on test set: %d %%' % (100 * correct / total))  # 正确率

运行一下测试代码也没什么问题。

然后是只保存模型的参数，这也是最常用的模型保存方法，只需要保存的时候传入模型的参数model.state_dict()即可，保存的模型文件后缀不变

1
2
# 仅保存模型参数
torch.save(model.state_dict(), 'model_dict.pt')

而加载模型时，由于仅保存了模型参数，因此需要先定义正确的模型结构，再加载模型参数

1
2
loaded_model_dict = nn.Sequential(nn.Flatten(), nn.Linear(784, 10))  # 定义模型结构
loaded_model_dict.load_state_dict(torch.load('model_dict.pt'))  # 读取模型参数
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
loaded_model_dict.eval()  # 调整为测试模式
correct = 0
total = 0

with torch.no_grad():  # 不需要计算梯度
    for data, target in MNIST_test:
        data, target = data.to(device), target.to(device)  # 将数据放到GPU上
            
        output = model(data)  # 计算输出
            
        predicted = torch.argmax(output, 1)  # 向量转换为标签
            
        total += target.size(0)  # 测试集数量
        correct += (predicted == target).sum().item()  # 预测值与正确类别是否相等，相等则+1（labels是类别数）
            
print('Accuracy on test set: %d %%' % (100 * correct / total))  # 正确率

Pytorch读取数据集

最后，我们来解决上文的遗留问题，如何读取自己的数据集，需要注意的是，由于深度学习涉及多种类型的数据，可能是图像，文字，音频，视频等，因此其读取的方式也不尽相同，下面只会介绍两种最通用的，适用于所有数据类型的数据集读取方式，本小节我们会以一个非常经典的猫狗数据集为例，这是一个包含25000张猫狗图像的二分类数据集，我们会以读取其训练集为目标。

Dataset类

首先，我们介绍Pytorch官方推荐的方法，即通过自定义类，配合DataLoader实现数据集的读取，DataLoader上文我们就讲过，其作用非常简单，输入一个Dataset类，返回一个python迭代器，用于实现数据集的小批量读取，而读取数据集的关键就在于Dataset类的实现上。

Dataset位于torch.utils.data，类似于nn.Module，Dataset通常作为父类被继承，并且在此基础上实现__init__()方法，__len__()方法与__getitem__()方法，即可被称为一个Dataset类，其中，__init__()方法用于初始化类，__len__()方法用于返回数据集的大小，而__getitem__()方法则根据索引，返回对应的样本与标签。

从上述对Dataset类的描述中不难看出，对于不同类型的数据，不同的存储方式，不同的目录结构，Dataset类会有不同的实现方式，只需要满足上述三个方法的要求即可，因此，这一部分的内容并不像上一小节的代码那么固定。

现在，我们就要定义一个Dataset类来读取猫狗数据集，该数据集下载下来是一个.zip压缩包，解压后是一个train文件夹，图像数据均在此文件夹中，数据均为.jpg格式，命名为类型加编号，例如cat.0.jpg，将该压缩包放置在./Data/CatDog_data，下面我们将采用不对数据进行提前解压，直接从压缩包中读取数据的方式，完整Dataset类定义如下

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
import zipfile  # 读取zip压缩包数据
from PIL import Image  # 读取图像数据
import os
import io

# 读取猫狗数据集
class CatDogDataset(Dataset):
    def __init__(self, root, transform=None, label_cat=0):
        self.root = root  # 数据集路径
        self.transform = transform
        self.label_cat = label_cat
        if self.label_cat == 0:
            self.label_dog = 1
        else:
            self.label_dog = 0
        
        with zipfile.ZipFile(self.root, 'r') as zip_ref:
            self.file_list = [f for f in zip_ref.namelist() if f.endswith('.jpg')]  # jpg文件列表
        
    def __getitem__(self, idx):
        with zipfile.ZipFile(self.root, 'r') as zip_ref:
            with zip_ref.open(self.file_list[idx]) as file:
                # 由于PIL读取图像数据是延迟加载的，此处用PIL直接读取会导致文件关闭了，但是图像数据还没有读取，导致报错
                img_data = file.read()  # 将文件完全读取至内存中
                
        # 使用io.BytesIO将二进制文件转化为数据流，保证不依赖zip文件
        img = Image.open(io.BytesIO(img_data)).convert('RGB')  # 保证读取图像为RGB三通道
        
        if self.transform is not None:
            img = self.transform(img)  # 对图像数据做变换
            
        label_name = os.path.basename(self.file_list[idx])  # 提取文件名
        if label_name.startswith('cat'):
            label = self.label_cat
        else:
            label = self.label_dog
                
        return img, label
        
    def __len__(self):  # 返回数据集大小
        return len(self.file_list)

现在我们来讲解上述代码，首先是导入的库，zipfile库用于打开zip文件，PIL用于读取图像数据，其余为辅助库。然后是__init__()方法，即类的初始化部分，该方法接受的参数有三个，root为数据集压缩包的路径，transform为需要对数据做的处理，因为torchvision中的transforms模块集成了许多用于处理图像的函数，且不同的图像预处理会影响训练结果，因此这部分通常会在类外进行设置，通过参数传入，label_cat为猫的类别编码，默认为0。在该方法中，除了初始化self变量外，我们还需要读取zip文件中的所有图像数据的列表，使用zipfile打开zip文件并使用.namelist()方法返回文件列表，需要注意的是该方法会返回文件夹路径，例如./train/，该路径我们是不需要的，我们只需要.jpg文件的路径，使用.endswith()方法检测路径是否以.jpg结尾进行筛选，得到文件路径列表self.file_list。

接着，我们先来看__len__()方法，该方法的任务是返回数据集大小，对于该数据集就是图像的数量，我们刚才已经拿到了文件列表，文件列表的长度就是图像的数量，因此直接return len(self.file_list)即可。

最后是__getitem__()方法，该方法的任务是接受一个索引，返回该索引对应的数据与标签即可，而文件列表刚好就可以实现索引与图像数据的对应，因此，实现思路为根据输入索引读取文件列表中对应的图像，并根据图像名称给出类别标签。

注：使用file.read()读取图像文件，而不是直接使用PIL读取，是因为PIL读取图像数据是后加载的，如果直接在with中使用PIL读取数据，会导致PIL在读取文件时，zip文件已经关闭，无法顺利读取而报错。

最后，我们将我们定义好的Dataset类实例化

1
2
3
4
5
6
catdog_path = '.\Data\CatDog_data\\train.zip'
transform_train = transforms.Compose([
    transforms.Resize((224, 224)),  # 数据集中图片大小不统一，需要统一转化为相同大小
    transforms.ToTensor()
])
data_catdog = CatDogDataset(catdog_path, transform=transform_train)

需要注意的是，在transforms中，相比于读取MNIST数据集，多了一步将所有的图像调整至统一大小。之后，再配合DataLoader实现小批量读取

1
catdog_train = DataLoader(data_catdog, batch_size=64, shuffle=True)
1
2
3
4
for data, target in catdog_train:
    print(data.shape)
    print(target.shape)
    break
1
2
torch.Size([64, 3, 224, 224])
torch.Size([64])

总结一下，使用自定义Dataset类读取数据是一种通用的数据读取方法，只需要按照定义实现指定的方法，就可以实现数据的读取，并且可以配合DataLoader实现数据的小批量读取，不过Dataset类的实现需要具体情况具体实现，并且其自由度并非最高，而且对于特定的数据，例如读取特定目录结构的图像数据，可以使用torchvision中的ImageFolder函数，其便捷程度是要更高的。

自定义迭代器

然后，我们来介绍一种自由度最高的数据读取方法，那就是完全手动读取，Pytorch训练时需要一个可以小批量读取数据的python迭代器，那么我们只要以写一个python迭代器为最终目标，其余的部分我们都可以自行定义与实现。其主要用于Dataset类配合DataLoader也无法实现全部需求的情况，例如文本数据，为了保证数据大小一致，通常会对文本数据进行填充，那么在训练时，除了需要数据本身，可能还需要文本数据的真实长度，此时由于DataLoader方法无法进行自定义，因此一般会选择手动读取数据。

现在，我们演示一下如何通过实现一个类来实现读取猫狗数据集，完整代码如下

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
import zipfile  # 读取zip压缩包数据
from PIL import Image  # 读取图像数据
import os
import io
import numpy as np

class CatDogData():
    def __init__(self, root, batch_size, transform=None, label_cat=0):
        self.root = root
        self.batch_size = batch_size
        self.transform = transform
        self.label_cat = label_cat
        if self.label_cat == 0:
            self.label_dog = 1
        else:
            self.label_dog = 0
        
        with zipfile.ZipFile(self.root, 'r') as zip_ref:
            self.file_list = [f for f in zip_ref.namelist() if f.endswith('.jpg')]  # jpg文件列表
            
        self.file_list = np.array(self.file_list)
        
    def __iter__(self):
        return self.__data_iter()
        
    def __data_iter(self):
        data_size = self.__len__()  # 数据集数量
        index = np.random.permutation(data_size)  # 生成随机索引
        
        for i in range(0, data_size, self.batch_size):
            batch_index = index[i: min(i + self.batch_size, data_size)]  # batch的索引列表
            batch_file_list = self.file_list[batch_index]
            
            with zipfile.ZipFile(self.root, 'r') as zip_ref:
                img_data_list = []
                for file_name in batch_file_list:
                    with zip_ref.open(file_name) as file:
                        img_data = file.read()  # 将文件完全读取至内存中
                        img_data_list.append(img_data)
                    
            X = []
            y = []
            for i in range(np.size(batch_file_list)):
                img = Image.open(io.BytesIO(img_data_list[i])).convert('RGB')  # 保证读取图像为RGB三通道
                if self.transform is not None:
                    img = self.transform(img)
                
                label_name = os.path.basename(batch_file_list[i])  # 提取文件名
                if label_name.startswith('cat'):
                    label = self.label_cat
                else:
                    label = self.label_dog
                    
                X.append(img)
                y.append(label)
                
            X = torch.stack(X, dim=0)
            y = torch.tensor(y)
            
            yield X, y
            
    def __len__(self):  # 返回数据集大小
        return np.size(self.file_list)

首先是__init__()方法，除了加入了batch_size参数以及将file_list转换为array之外，与使用Dataset类的实现没有区别。然后是该类的实现重点，__iter__()方法与__data_iter函数，这里采用了与传统的python迭代器不同的实现方法，没有实现__next__()方法，而是使用yield关键字实现python迭代器，再使用__iter__()方法返回该迭代器，之所以必须实现__iter__()方法，是因为如果没有实现该方法，则该类实例化之后无法被识别为一个可迭代对象，也因此无法配合for循环实现小批量读取。而__data_iter函数中，将图像数据的读取与小批量读取同时进行实现，其中图像数据的读取代码跟Dataset类的实现是一致的，而实现小批量读取的代码可以参考本站的另外一篇博客Python：迭代器与生成器 。

最后，将定义好的类实例化

1
2
3
4
5
6
catdog_path = 'E:\Code\DL\Data\CatDog_data\\train.zip'
transform_train = transforms.Compose([
    transforms.Resize((224, 224)),  # 数据集中图片大小不统一，需要统一转化为相同大小
    transforms.ToTensor()
])
data_catdog = CatDogData(catdog_path, batch_size=64, transform=transform_train)
1
2
3
4
for data, target in data_catdog:
    print(data.shape)
    print(target.shape)
    break
1
2
torch.Size([64, 3, 224, 224])
torch.Size([64])