小土堆 Dataset

最新推荐文章于 2024-09-10 13:33:13 发布

简.454

最新推荐文章于 2024-09-10 13:33:13 发布

阅读量337

点赞数

CC 4.0 BY-SA版权

文章标签： python 深度学习开发语言

本文链接：https://blog.csdn.net/weixin_62655710/article/details/127128293

该博客展示了如何使用PyTorch的`torch.utils.data.Dataset`类创建自定义数据集。通过定义`__init__`和`__getitem__`方法，可以轻松获取每个数据样本及其对应的标签。同时，`__len__`方法用于获取数据集的长度。示例中创建了两个数据集分别对应ants和bees，并将它们合并为训练数据集。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

Dataset类：提供了一种方式去获取数据及其label

主要应用：
- 如何获取每一个数据和对应的label
- 如何获取数据的长度

导入

from torch.utils.data import Dataset

os.path.join 对路径进行拼接。
os.listdir() 返回指定的文件夹包含的文件名字的列表。


#展示Dataset类的用法：提供了一种获取数据的和对应label的方式
""""
    1.如何获取每一个数据及其label
    2.获取数据的长度
"""
from torch.utils.data import Dataset  #引入对应的package
import os
from PIL import Image #图片操作的package

class MyData(Dataset):#建立一个类，要继承Dataset类
#root_dir是train的相对地址 label_dir是ants
    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.label_dir,img_name)
        img=Image.open(img_item_path)
        label=self.label_dir
        return img,label
    def __len__(self):
        return len(self.img_path)

root_dir="hymenoptera_data/train"
ants_label_dir="ants"
bees_label_dir="bees"
ants_dataset=MyData(root_dir,ants_label_dir)
bees_dataset=MyData(root_dir,bees_label_dir)
train_dataset=ants_dataset+bees_dataset
img,label=ants_dataset[0]
#img.show()
print(len(train_dataset))