PyTorch/[PyTorch 学习笔记] 2.1 DataLoader 与 DataSet
本章代码:https://github.com/zhangxiann/PyTorch_Practice/blob/master/lesson2/rmb_classification/
人民币 二分类
实现 1 元人民币和 100 元人民币的图片二分类。前面讲过 PyTorch 的五大模块:数据、模型、损失函数、优化器和迭代训练。
数据模块又可以细分为 4 个部分:
- 数据收集:样本和标签。
- 数据划分:训练集、验证集和测试集
- 数据读取:对应于PyTorch 的 DataLoader。其中 DataLoader 包括 Sampler 和 DataSet。Sampler 的功能是生成索引, DataSet 是根据生成的索引读取样本以及标签。
- 数据预处理:对应于 PyTorch 的 transforms
# DataLoader 与 DataSet
torch.utils.data.DataLoader()
1 | torch.utils.data.DataLoader(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 类,决定数据从哪里读取以及如何读取
- batchsize: 批大小
- num_works:num_works: 是否多进程读取数据
- sheuffle: 每个 epoch 是否乱序
- drop_last: 当样本数不能被 batchsize 整除时,是否舍弃最后一批数据
Epoch, Iteration, Batchsize
- Epoch: 所有训练样本都已经输入到模型中,称为一个 Epoch
- Iteration: 一批样本输入到模型中,称为一个 Iteration
- Batchsize: 批大小,决定一个 iteration 有多少样本,也决定了一个 Epoch 有多少个 Iteration
假设样本总数有 80,设置 Batchsize 为 8,则共有 \(80 \div 8=10\) 个 Iteration。这里 \(1 Epoch = 10 Iteration\)。
假设样本总数有 86,设置 Batchsize 为 8。如果drop_last=True
则共有 10 个 Iteration;如果drop_last=False
则共有 11 个 Iteration。
torch.utils.data.Dataset
功能:Dataset 是抽象类,所有自定义的 Dataset 都需要继承该类,并且重写__getitem()__
方法和__len__()
方法 。__getitem()__
方法的作用是接收一个索引,返回索引对应的样本和标签,这是我们自己需要实现的逻辑。__len__()
方法是返回所有样本的数量。
数据读取包含 3 个方面
- 读取哪些数据:每个 Iteration 读取一个 Batchsize 大小的数据,每个 Iteration 应该读取哪些数据。
- 从哪里读取数据:如何找到硬盘中的数据,应该在哪里设置文件路径参数
- 如何读取数据:不同的文件需要使用不同的读取方法和库。
这里的路径结构如下,有两类人民币图片:1 元和 100 元,每一类各有 100 张图片。
- RMB_data
- 1
- 100
首先划分数据集为训练集、验证集和测试集,比例为 8:1:1。
数据划分好后的路径构造如下:
- rmb_split
- train
- 1
- 100
- valid
- 1
- 100
- test
- 1
- 100
- train
实现读取数据的 Dataset,编写一个get_img_info()
方法,读取每一个图片的路径和对应的标签,组成一个元组,再把所有的元组作为 list 存放到self.data_info
变量中,这里需要注意的是标签需要映射到 0 开始的整数: rmb_label = {"1": 0, "100": 1}
。
1 | @staticmethod |
然后在Dataset
的初始化函数中调用get_img_info()
方法。
1 | def __init__(self, data_dir, transform=None): |
然后在__getitem__()
方法中根据index
读取self.data_info
中路径对应的数据,并在这里做 transform 操作,返回的是样本和标签。
1 | def __getitem__(self, index): |
在__len__()
方法中返回self.data_info
的长度,即为所有样本的数量。
1 | # 返回所有样本的数量 |
在train_lenet.py
中,分 5 步构建模型。
第 1 步设置数据。首先定义训练集、验证集、测试集的路径,定义训练集和测试集的transforms
。然后构建训练集和验证集的RMBDataset
对象,把对应的路径和transforms
传进去。再构建DataLoder
,设置 batch_size,其中训练集设置shuffle=True
,表示每个 Epoch 都打乱样本。
1 | # 构建MyDataset实例train_data = RMBDataset(data_dir=train_dir, transform=train_transform)valid_data = RMBDataset(data_dir=valid_dir, transform=valid_transform) |
第 2 步构建模型,这里采用经典的 Lenet 图片分类网络。
1 | net = LeNet(classes=2) |
第 3 步设置损失函数,这里使用交叉熵损失函数。
1 | criterion = nn.CrossEntropyLoss() |
第 4 步设置优化器。这里采用 SGD 优化器。
1 | optimizer = optim.SGD(net.parameters(), lr=LR, momentum=0.9) # 选择优化器 |
第 5 步迭代训练模型,在每一个 epoch 里面,需要遍历 train_loader 取出数据,每次取得数据是一个 batchsize 大小。这里又分为 4 步。第 1 步进行前向传播,第 2 步进行反向传播求导,第 3 步使用optimizer
更新权重,第 4 步统计训练情况。每一个 epoch 完成时都需要使用scheduler
更新学习率,和计算验证集的准确率、loss。
1 | for epoch in range(MAX_EPOCH): |
我们可以看到每个 iteration,我们是从train_loader
中取出数据的。
1 | def __iter__(self): |
这里我们没有设置多进程,会执行_SingleProcessDataLoaderIter
的方法。我们以_SingleProcessDataLoaderIter
为例。在_SingleProcessDataLoaderIter
里只有一个方法_next_data()
,如下:
1 | def _next_data(self): |
在该方法中,self._next_index()
是获取一个 batchsize 大小的 index 列表,代码如下:
1 | def _next_index(self): |
其中调用的sampler
类的__iter__()
方法返回 batch_size 大小的随机 index 列表。
1 | def __iter__(self): |
然后再返回看 dataloader
的_next_data()
方法:
1 | def _next_data(self): |
在第二行中调用了self._dataset_fetcher.fetch(index)
获取数据。这里会调用_MapDatasetFetcher
中的fetch()
函数:
1 | def fetch(self, possibly_batched_index): |
这里调用了self.dataset[idx]
,这个函数会调用dataset.__getitem__()
方法获取具体的数据,所以__getitem__()
方法是我们必须实现的。我们拿到的data
是一个 list,每个元素是一个 tunple,每个 tunple 包括样本和标签。所以最后要使用self.collate_fn(data)
把 data 转换为两个 list,第一个 元素 是样本的batch 形式,形状为 [16, 3, 32, 32] (16 是 batch size,[3, 32, 32] 是图片像素);第二个元素是标签的 batch 形式,形状为 [16]。
所以在代码中,我们使用inputs, labels = data
来接收数据。
PyTorch 数据读取流程图
首先在 for 循环中遍历DataLoader
,然后根据是否采用多进程,决定使用单进程或者多进程的DataLoaderIter
。在DataLoaderIter
里调用Sampler
生成Index
的 list,再调用DatasetFetcher
根据index
获取数据。在DatasetFetcher
里会调用Dataset
的__getitem__()
方法获取真正的数据。这里获取的数据是一个 list,其中每个元素是 (img, label) 的元组,再使用 collate_fn()
函数整理成一个 list,里面包含两个元素,分别是 img 和 label 的tenser
。
下图是我们的训练过程的 loss 曲线:
参考资料
如果你觉得这篇文章对你有帮助,不妨点个赞,让我有更多动力写出好文章。
我的文章会首发在公众号上,欢迎扫码关注我的公众号张贤同学。