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

一、自定义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