Pytorch学习笔记一：制作、加载自己的图像数据集
文章目录
前言
首先介绍如何用pytorch加载网络现有数据集，然后介绍如何制作自己的图像数据集并批量读取来训练自己的网络。
提示：以下是本篇文章正文内容，下面案例可供参考
一、下载数据集
使用Pytorch进行读取本地的MINIST数据集并进行装载
# 训练数据和测试数据的下载
trainDataset = torchvision.datasets.MNIST( # torchvision可以实现数据集的训练集和测试集的下载
root="./data", # 下载数据，并且存放在data文件夹中
train=True, # train用于指定在数据集下载完成后需要载入哪部分数据，如果设置为True，则说明载入的是该数据集的训练集部分；如果设置为False，则说明载入的是该数据集的测试集部分。
transform=transforms.ToTensor(), # 数据的标准化等操作都在transforms中，此处是转换
download=True
)
testDataset = torchvision.datasets.MNIST(
root="./data",
train=False,
transform=transforms.ToTensor(),
download=True
)
二、加载自己的数据集
1.制作数据集
训练神经网络需要标准输入图像和它的真值标签。
在分类问题中，比如猫、狗、船、车等等，我们可以用数字代表不同的分类。可以制作一个txt文档用于存放输入图像的地址和它对应的标签数字。
我现在有个任务需要以图像作为输入，以另一张处理过后的图像作为它的真值，所以我在txt文本下面写的是它们的路径。在项目路径下新建了一个train文件夹用于放训练图片，并在train文件夹下新建一个训练的txt用于标注训练图像和标签图像
2.加载数据集
Dataset类
PyTorch读取图片，主要是通过Dataset类是Pytorch中所有数据集加载类中应该继承的父类。我们通过继承改写Dataset类来读取自己的图像数据集。其中以下三个函数必须改写：
__init__方法里面进行读取数据文件
__getitem__方法进行支持下标访问
__len__方法返回自定义数据集的大小，方便后期遍历
class OpticalSARDataset(Data.Dataset):
"""
定义自己的数据集、读取数据、初始化数据
"""
def __init__(self, data_dir, part):
# 所有图片的绝对路径
assert part in ["train", "val"]
self.image_dir = os.path.join(data_dir, part)
self.img_names = []
self.label_names = []
with open(os.path.join(data_dir, part, "label.txt")) as f:
while True:
il = f.readline(1500) # 如果样本数据名称大于1500，修改该值
if not il:
break
a = il.split(sep=' ')
self.img_names.append(a[0])
self.label_names.append(a[1][0:-1]) # remove '\n'
self.samples_num = len(self.img_names)
# print(self.samples_num)
self.transform = torchvision.transforms.Compose([
# 将 PIL 图片转换成位于[0.0, 1.0]的floatTensor, shape (C x H x W)
torchvision.transforms.ToTensor()])
def __len__(self):
# 返回图像的数量
return self.samples_num
def __getitem__(self, index):
tp_img = Image.open(os.path.join(self.image_dir, self.img_names[index])
).convert('RGB')
tp_label = Image.open(os.path.join(self.image_dir, self.label_names[index])
).convert('RGB')
# PIL.Image.open 读取的图片数据是RGB格式；
tp_img = cv2.cvtColor(np.asarray(tp_img), cv2.COLOR_RGB2BGR)
tp_label = cv2.cvtColor(np.asarray(tp_label), cv2.COLOR_RGB2BGR) # 转换为BGR便于cv2.imshow，跟下面imshow之前RGB2BGR只用一种方法，这里统一为cv2的BGR格式
img = self.transform(tp_img)
label = self.transform(tp_label)
sample = {
"label": label, # shape
"image": img # shape: (3, *image_size)
}
return sample
定义数据集
# 利用之前创建好的OpticalSARDataset类去创建数据对象
train_dataset = OpticalSARDataset(data_dir, 'train') # 训练数据集
Dataloader类
之前所说的Dataset类是读入数据集数据并且对读入的数据进行了索引。
但是光有这个功能是不够用的，在实际的加载数据集的过程中，我们的数据量往往都很大，对此我们还需要一下几个功能：
可以分批次读取：batch-size
可以对数据进行随机读取，可以对数据进行洗牌操作(shuffling)，打乱数据集内数据分布的顺序
可以并行加载数据(利用多核处理器加快载入数据的效率)
Dataloader这个类并不需要我们自己设计代码，我们只需要利用DataLoader类读取我们设计好的类。
实例化数据集
# 利用dataloader读取我们的数据对象，并设定batch-size和工作现场
train_iter = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size,
shuffle=True, num_workers=0)
batch = iter(train_iter).next()
print(batch["image"].shape, batch["label"].shape)
print(batch["image"][0].shape)
总结
参考博客：
定义自己的数据集
pytorch加载自己数据集
设计自己的数据
训练自己数据完整步骤
Dataset类