【Pytorch入门学习】 Dateset类的实现笔记（套路） 附有详细的代码注释以及实验框架流程

---

1.项目目录

本实验使用的数据集是蚂蚁蜜蜂数据集蚂蚁蜜蜂数据集ants & bees

下载完成后，对数据集进行解压，整个项目的组织结构如下所示

P7_dataset1.py是我们本次实验的python脚本

打开ants目录下面就是所有的文件名

---

2.实现框架

Dataset中主要有三个需要我们重写的函数，__init__(), __getitem__(), __len__() ，这个也是面试的时候经常被问的问题，下面讲解主要的实现步骤

1. 首先定义自己的一个类，继承自Dataset
2. 实现__init__(self, root_dir, label_dir)函数，分别传进数据的目录以及数据标签（主要用于实现对数据集的定位）
3. 实现__getitem__(self, item)传进用户的定位标签,用户可以通过使用dataset[item]来直接从对应数据集标签目录中获取想要的数据
4. 实现 __len__(self)主要返回数据集的长度
5. dataset之间的相加可以直接用 “+” 号相加，获取一个更大的数据集

---

3. 实现代码

主要以获取图片为例子，有详细注释

# -*- coding:utf-8 -8-
"""
Author: Leung
Date: 2022--10--27
"""
from torch.utils.data import Dataset
from PIL import Image
import os


class MyData(Dataset):
    def __init__(self, root_dir, label_dir):
        # 传入图片文件夹，该文件夹名字对应的就是这类数据集的label
        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)   # 将路径名称变为一个列表，方便后面getitem函数进行索引操作
                                                # 该列表存放的是整个目录下面所有的文件名

    def __getitem__(self, item):
        img_name = self.img_path[item]      # 利用上面初始化函数建立的函数
        img_path_name = os.path.join(self.path, img_name)   # 同上，目录相加
        img = Image.open(img_path_name)     # 打开文件图片，方便后续返回
        label = self.label_dir              # 每个图片的标签
        return img, label                   # 返回图片 + 标签的形式

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


root_dir = "data/train"
ants_label_dir = "ants"
bees_label_dir = "bees"
ants_dataset = MyData(root_dir, label_dir=ants_label_dir)       # 蚂蚁数据集
bees_dataset = MyData(root_dir, label_dir=bees_label_dir)       # 蜜蜂数据集

train_dataset = ants_dataset + bees_dataset # 直接相加，123是蚂蚁图片，124是蜜蜂图片

# 从创建的实例中获取数据
img123_ants, label_123 = train_dataset[123]
img124_bees, label_124 = train_dataset[124]

# 展示图片
img123_ants.show()
img124_bees.show()

程序运行结果

红字为注释

---

写在最后

各位看官，都看到这里了，麻烦动动手指头给博主来个点赞8，您的支持作者最大的创作动力哟！
注：官方文档说了很多，网上也有很多精彩的教程，我希望结合土堆的教学视频PyTorch深度学习快速入门教程谈一下自己对Dataset类的理解
才疏学浅，若有纰漏，恳请斧正
本文章仅用于各位作为学习交流之用，不作任何商业用途，若涉及版权问题请速与作者联系，望悉知