PyTorch重写DataSet类

PyTorch重写DataSet类


文章目录


前言

在之前沐神的Cifar-10分类 课程学习中,沐神是用的将每一类创建一个文件夹去完成图片的导入。此外我们还可以通过重写DataSet类来完成!

一、如何重写?

通过查看官方文档我们可知。

需要去重写__getitem__这个方法,去以一种特定的方法拿到一个数据。并且选择性的重写__len__这个方法,去返回整个数据集的大小。

二、具体代码

1.数据集格式

这个数据集是沐神课程上讲过的cifar-10数据集。

train和test文件夹分别为要进行训练和测试的图片。而训练数据的标签以csv文件存在trainLabels.csv文件中。

2.获取标签

python 复制代码
def read_csv_labels(fname):
    with open(fname,'r') as f:
        lines = f.readlines()[1:]
    tokens = [l.rstrip().split(',') for l in lines]
    return dict(((name,label) for name,label in tokens))

这里通过一个read_csv_labels的方法 将图片名字和标签以一个字典的方式返回

3.重写dataset

python 复制代码
class MyDateset(Dataset):
    def __init__(self,root_dir,state,label_dict=None):
        self.root_dir = root_dir
        self.state = state
        if label_dict is not None:
            self.label_dict = label_dict
        self.img_path = os.listdir(os.path.join(root_dir,state))
        # os.listdir 将当前文件夹下的图片名称按列表返回

    def __getitem__(self, idx):
        img = Image.open(os.path.join(self.root_dir,self.state,self.img_path[idx]))
        if self.state == 'train':
            img_num =self.img_path[idx].split('.')[0]
            # 这个取出来是数字.jpg 所以需要将.jpg舍去
            label = self.label_dict[img_num]
            return img,label
        else:
            return img

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

state参数表示此时是训练数据集还是测试数据集。

4.调用

python 复制代码
root_dir = "D:\\PytorchLearn\\cifar-10"
label_dict = read_csv_labels(os.path.join(root_dir,"trainLabels.csv"))

train_dataset = MyDateset(root_dir,'train',label_dict)

test_dataset = MyDateset(root_dir,'test')

train_iter = torch.utils.data.DataLoader(train_dataset,batch_size=8,shuffle=True)

总结

以上就是重写DataSet的方法,有不足之处还望各位指出。

相关推荐
学术小八1 小时前
2025年人工智能、虚拟现实与交互设计国际学术会议
人工智能·交互·vr
nbsaas-boot2 小时前
Java 正则表达式白皮书:语法详解、工程实践与常用表达式库
开发语言·python·mysql
仗剑_走天涯2 小时前
基于pytorch.nn模块实现线性模型
人工智能·pytorch·python·深度学习
chao_7892 小时前
二分查找篇——搜索旋转排序数组【LeetCode】两次二分查找
开发语言·数据结构·python·算法·leetcode
cnbestec3 小时前
协作机器人UR7e与UR12e:轻量化设计与高负载能力助力“小而美”智造升级
人工智能·机器人·协作机器人·ur协作机器人·ur7e·ur12e
zskj_zhyl3 小时前
毫米波雷达守护银发安全:七彩喜跌倒检测仪重构居家养老防线
人工智能·安全·重构
gaosushexiangji4 小时前
利用sCMOS科学相机测量激光散射强度
大数据·人工智能·数码相机·计算机视觉
ai小鬼头5 小时前
AIStarter新版重磅来袭!永久订阅限时福利抢先看
人工智能·开源·github
说私域6 小时前
从品牌附庸到自我表达:定制开发开源AI智能名片S2B2C商城小程序赋能下的营销变革
人工智能·小程序
飞哥数智坊6 小时前
新版定价不够用,Cursor如何退回旧版定价
人工智能·cursor