2.pytorch学习:自己导入数据集与利用自带的公开数据集

最新推荐文章于 2024-10-07 21:08:20 发布

原创

最新推荐文章于 2024-10-07 21:08:20 发布 · 2.5k 阅读

31 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch

本文介绍如何创建自定义数据集类并集成到PyTorch中，包括数据集的初始化、索引和长度定义。此外，还介绍了如何使用PyTorch自带的数据集，并演示了数据预处理步骤。

自己导入数据集

蚂蚁蜜蜂分类数据集下载链接：

https://download.pytorch.org/tutorial/hymenoptera_data.zip

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


class MyData(Dataset):
    # root_dir数据集根目录文件夹，label_dir为标注过的图片文件夹
    # 初始化+读取数据集
    def __init__(self, root_dir, label_dir):
        # root_dir数据集根目录文件夹
        self.root_dir = root_dir
        # label_dir为标注过的图片文件夹
        self.label_dir = label_dir
        # 使用os.path.join()函数，拼接路径，因为win是\\拼接，linux是\拼接。
        self.path = os.path.join(self.root_dir, self.label_dir)
        # 将路径下的文件存成数组(array)的形式。数组的元素对应每个图片的名字(str字符串类型)。
        self.img_path = os.listdir(self.path)

    # 对于指定的idx(索引，因为img_path是一个由图片名（字符串）按照一定顺序组成的数组)，获取数据并返回。
    def __getitem__(self, idx):
        # 通过idx索引图片名。
        img_name = self.img_path[idx]
        # 拼接数据集根目录、标注图片文件夹目录与图片名。
        img_item_path = os.path.join(self.root_dir, self.label_dir, img_name)
        # 利用PIL(Python Imaging Library)中的Image中的open()函数打开图片。
        img = Image.open(img_item_path)
        # 因为这是分类任务，故标注图片的文件夹名就是种类的名字。
        label = self.label_dir
        # 返回图片与
        return img, label

    # 返回图片的个数(图片文件名列表的长度)。
    def __len__(self):
        return len(self.img_path)


# 数据集根目录。
dataset_root_dir = "hymenoptera_data/hymenoptera_data/train"
# 数据集的标注文件夹，因为是蚂蚁的蜜蜂的分类问题。
ants_label_dir = "ants"
bees_label_dir = "bees"
# def __init__(self, root_dir, label_dir): ,方法__init__有两个形参,根据这个类创建实例就必须要指定形参的值。
ants_dataset = MyData(dataset_root_dir, ants_label_dir)
bees_dataset = MyData(dataset_root_dir, ants_label_dir)
# 将两个数据集加载一块组成训练集
train_dataset = ants_dataset + bees_dataset

注意的点1：

方法__init__有两个形参,根据这个类创建实例就必须要指定形参的值。这个目的是将两个路径通过os.path.join()函数，拼成想要