Pytorch框架学习记录1——Dataset类代码实战

介绍

• torch.utils.data.Dataset是代表自定义数据集方法的抽象类，你可以自己定义你的数据类继承这个抽象类，非常简单，只需要定义__len__和__getitem__这两个方法就可以。
• 通过torch.utils.data.DataLoader类来定义一个新的迭代器，用来将自定义的数据读取接口的输出或者PyTorch已有的数据读取接口的输入按照batch size封装成Tensor，后续只需要再包装成Variable即可作为模型的输入。

这两个抽象类中用到的python知识点

能够熟练的使用python语言的技巧，是理解pytorch源码的关键。在torch.utils.data.Dataset和torch.utils.data.DataLoader这两个类中会用到python抽象类的魔法方法，包括__len__(self)，getitem(self)和__iter__(self)

• __len__(self) 定义当被len()函数调用时的行为（返回容器中元素的个数）
• __getitem__(self)定义获取容器中指定元素的行为，相当于self[key]，即允许类对象可以有索引操作。
• __iter__(self)定义当迭代容器中的元素的行为

数据集下载地址：https://pan.baidu.com/s/1qNCOVz15mCSQEDZZXJAaoQ?pwd=qz2b
提取码：qz2b

1. 导入包

from torch.utils.data import Dataset
from PIL import Image
import os

2. 创建类

创建子类MyData，继承父类Dataset，并对函数进行重写。

class MyData(Dataset):

    def __init__(self, root_dir, label_dir):
        self.root_dir = root_dir
        self.label_dir = label_dir
        self.path = os.path.join(self.root_dir, self.label_dir)  # 路径拼接
        self.img_path = os.listdir(self.path)  # 获得文件夹下所有的文件名

    def __getitem__(self, idx):
        img_name = self.img_path[idx]
        img_item_path = os.path.join(self.root_dir, self.img_path, img_name)
        img = Image.open(img_item_path)
        label = self.label_dir
        return img, label

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

3. 调用

类进行实例化，并进行调用

root_dir = "C:\\Users\\hp\\PycharmProjects\\pythonProject\\Pytorch_Learning\\flower_data\\train"
daisy_label_dir = "daisy"
roses_label_dir = "roses"
daisy_dataset = MyData(root_dir, daisy_label_dir)
roses_dataset = MyData(root_dir, roses_label_dir)

train_dataset = daisy_dataset + roses_dataset
print("daisy:",len(daisy_dataset),"\nroses",len(roses_dataset),"\ndaisy+roses",len(train_dataset))
img, label = train_dataset[0]
img2, label = train_dataset[577]
img2.show()