PyTorch 学习笔记（一）：让PyTorch读取你的数据集

本文截取自《PyTorch 模型训练实用教程》，获取全文pdf请点击：https://github.com/tensor-yu/PyTorch_Tutorial

---

号外：20万字的《PyTorch实用教程》（第二版）于 2024 年 4 月开源了！
号外：20万字的《PyTorch实用教程》（第二版）于 2024 年 4 月开源了！
号外：20万字的《PyTorch实用教程》（第二版）于 2024 年 4 月开源了！

---

在第一版的基础上，融入八个计算机视觉实战项目、五个自然语言处理应用，以及四个前沿大语言模型部署实例。此外，针对技术落地的关键环节，还深入讲解了 ONNX 与 TensorRT 的应用，铺平工程部署的道路。
愿 20 万字的教程助力每位读者更上一层楼！
敬请关注：https://github.com/TingsongYu/PyTorch-Tutorial-2nd

想让PyTorch能读取我们自己的数据，首先要了解pytroch读取图片的机制和流程，然后按流程编写代码。

Dataset类

PyTorch读取图片，主要是通过Dataset类，所以先简单了解一下Dataset类。Dataset类作为所有的datasets的基类存在，所有的datasets都需要继承它，类似于C++中的虚基类。

源码如下：

class Dataset(object):
"""An abstract class representing a Dataset.
All other datasets should subclass it. All subclasses should override
``__len__``, that provides the size of the dataset, and ``__getitem__``,
supporting integer indexing in range from 0 to len(self) exclusive.
"""
def __getitem__(self, index):
	raise NotImplementedError
def __len__(self):
	raise NotImplementedError
def __add__(self, other):
	return ConcatDataset([self, other])

这里重点看 getitem函数，getitem接收一个index，然后返回图片数据和标签，这个index通常指的是一个list的index，这个list的每个元素就包含了图片数据的路径和标签信息。

然而，如何制作这个list呢，通常的方法是将图片的路径和标签信息存储在一个txt中，然后从该txt中读取。
那么读取自己数据的基本流程就是：

1. 制作存储了图片的路径和标签信息的txt
2. 将这些信息转化为list，该list每一个元素对应一个样本
3. 通过getitem函数，读取数据和标签，并返回数据和标签

在训练代码里是感觉不到这些操作的，只会看到通过DataLoader就可以获取一个batch的数据，其实触发去读取图片这些操作的是DataLoader里的__iter__(self)，后面会详细讲解读取过程。在本小节，主要讲Dataset子类。
因此，要让PyTorch能读取自己的数据集，只需要两步：

1. 制作图片数据的索引

2. 构建Dataset子类

3. 制作图片数据的索引
   这个比较简单，就是读取图片路径，标签，保存到txt文件中，这里注意格式就好
   特别注意的是，txt中的路径，是以训练时的那个py文件所在的目录为工作目录，所以这里需要提前算好相对路径！
   运行代码 Code/1_data_prepare/1_3_generate_txt.py，即会在/Data/文件夹下面看到 train.txt valid.txt
   txt中是这样的：
   

构建Dataset子类

下面是本实验构建的Dataset子类——MyDataset类：

# coding: utf-8
from PIL import Image
from torch.utils.data import Dataset
class MyDataset(Dataset):
def __init__(self, txt_path, transform = None, target_transform = None):
	fh = open(txt_path, 'r')
	imgs = []
	for line in fh:
		line = line.rstrip()
		words = line.split()
		imgs.append((words[0], int(words[1])))
		self.imgs = imgs 
		self.transform = transform
		self.target_transform = target_transform
def __getitem__(self, index):
	fn, label = self.imgs[index]
	img = Image.open(fn).convert('RGB') 
	if self.transform is not None:
		img = self.transform(img) 
	return img, label
def __len__(self):
	return len(self.imgs)

首先看看初始化，初始化中从我们准备好的txt里获取图片的路径和标签，并且存储在self.imgs，self.imgs就是上面提到的list，其一个元素对应一个样本的路径和标签，其实就是txt中的一行。

初始化中还会初始化transform，transform是一个Compose类型，里边有一个list，list中就会定义了各种对图像进行处理的操作，可以设置减均值，除标准差，随机裁剪，旋转，翻转，仿射变换等操作。

在这里我们可以知道，一张图片读取进来之后，会经过数据处理（数据增强），最终变成输入模型的数据。这里就有一点需要注意，PyTorch的数据增强是将原始图片进行了处理，并不会生成新的一份图片，而是“覆盖”原图，当采用randomcrop之类的随机操作时，每个epoch输入进来的图片几乎不会是一模一样的，这达到了样本多样性的功能。

然后看看核心的 getitem函数：

第一行：self.imgs 是一个list，也就是一开始提到的list，self.imgs的一个元素是一个元组，包含图片路径，图片标签，这些信息是从txt文件中读取

第二行：利用Image.open对图片进行读取，img类型为 Image ，mode=‘RGB’

第三行与第四行： 对图片进行处理，这个transform里边可以实现 减均值，除标准差，随机裁剪，旋转，翻转，放射变换，等等操作，这个放在后面会详细讲解。

当Mydataset构建好，剩下的操作就交给DataLoder，在DataLoder中，会触发Mydataset中的getiterm函数读取一张图片的数据和标签，并拼接成一个batch返回，作为模型真正的输入。下一小节将会通过一个小例子，介绍DataLoder是如何获取一个batch，以及一张图片是如何被PyTorch读取，最终变为模型的输入的。

转载请注明出处：https://blog.csdn.net/u011995719/article/details/85102770