PyTorch数据加载：自定义数据集【Dataset：处理每个原始样本】【DataLoader：每次生成batch_size个样本】【collate_fn：重新设置一个Batch中所有样本的加载格式】

系统管理员 2023-09-30 09:08 4阅读 0赞

## 一、自定义Dataset ##

Dataset是一个包装类：

*  用来将数据包装为Dataset类，然后传入DataLoader中，我们再使用DataLoader这个类来更加快捷的对数据进行操作。
 *  可以通过继承Dataset来将数据集的源文件、规模和其他非必要的功能打包，从而供DataLoader使用。

### 1、“文本分类”任务下使用自定义Dataset ###

class.txt：所有类别

finance
    realty
    stocks
    education
    science
    society
    politics
    sports
    game
    entertainment

train.txt：训练数据样式

中华女子学院：本科层次仅1专业招男生	3
    两天价网站背后重重迷雾：做个网站究竟要多少钱	4
    东5环海棠公社230-290平2居准现房98折优惠	1
    卡佩罗：告诉你德国脚生猛的原因 不希望英德战踢点球	7
    82岁老太为学生做饭扫地44年获授港大荣誉院士	5
    记者回访地震中可乐男孩：将受邀赴美国参观	5
    冯德伦徐若�隔空传情 默认其是女友	9
    传郭晶晶欲落户香港战伦敦奥运 装修别墅当婚房	1
    《赤壁OL》攻城战诸侯战硝烟又起	8
    “手机钱包”亮相科博会	4
    上海2010上半年四六级考试报名4月8日前完成	3
    李永波称李宗伟难阻林丹取胜 透露谢杏芳有望出战	7
    3岁女童下体红肿 自称被幼儿园老师用尺子捅伤	5
    金证顾问：过山车行情意味着什么	2
    谁料地王如此虚	1
    《光环5》Logo泄露 Kinect版几无悬念	8
    海淀区领秀新硅谷宽景大宅预计10月底开盘	1
    柴志坤：土地供应量不断从紧 地价难现07水平(图)	1
    伊达传说EDDA Online	8
    三联书店建起书香巷	4
    宇航员尿液堵塞国际空间站水循环系统	4
    研究发现开车技术差或与基因相关	6
    皇马输球替补席闹丑闻 队副女球迷公然调情(视频)	7
    北京建工与市政府再度合作推出郭庄子限价房	1
    组图：李欣汝素颜出镜拍低碳环保大片	9
    2008中文网志年会演讲人：庄秀丽	4
    3000点之下是买入好时机	2

dataset.py文件：

import torch
    from tqdm import tqdm
    from torch.utils.data import Dataset
    
    
    class SampleDataset(Dataset):
        def __init__(self, config, split):
            if split == 'train':
                self.texts, self.labels = read_from_txt(config.train_path)
            if split == 'val':
                self.texts, self.labels = read_from_txt(config.val_path)
            if split == 'test':
                self.texts, self.labels = read_from_txt(config.test_path)
    
            self.labels = torch.LongTensor(self.labels).to(config.device)
    
        def __len__(self):
            return len(self.labels)
    
        def __getitem__(self, idx):
            text = self.texts[idx]
            label = self.labels[idx]
    
            return text, label
    
    
    def