如果是可以一次性加载进内存的数据,上一篇博客:pytorch 构造读取数据的工具类 Dataset 与 DataLoader (pytorch Data学习一),已经足以应付了,但是很多时候数据集较大,比如6个T…的数据,没办法直接加载,因此这篇博客先学习构建一个可迭代的数据Dataset

MyIterableDataset构造

1. 迭代读取文本文件

要借助pytorch的IterableDataset模块,官方文档是:https://pytorch.org/docs/stable/data.html#torch.utils.data.IterableDataset,官方文档里侧重于对多进程时的描述,按照官网的说法,需要继承这个IterableDataset类,然后覆写__iter__这个方法,返回一个可迭代的对象即可。

因此构造一个基本类

from torch.utils.data import IterableDataset


class MyIterableDataset(IterableDataset):

    def __init__(self, file_path):
        self.file_path = file_path

    def __iter__(self):
        with open(self.file_path, 'r') as file_obj:
            for line in file_obj: # 更多操作在这里完成
                line_data = line.strip('\n').split(',')
                yield line_data

使用的时候直接调用即可:

if __name__ == '__main__':
    dataset = MyIterableDataset('test_csv.csv')
    for data in dataset:
        print(data)

完整代码为:

from torch.utils.data import IterableDataset


class MyIterableDataset(IterableDataset):

    def __init__(self, file_path):
        self.file_path = file_path

    def __iter__(self):
        with open(self.file_path, 'r') as file_obj:
            for line in file_obj:
                line_data = line.strip('\n').split(',')
                yield line_data


if __name__ == '__main__':
    dataset = MyIterableDataset('test_csv.csv')
    for data in dataset:
        print(data)

2.借用pandas迭代读取csv/excel等文件

改写__iter__方法即可:

from torch.utils.data import IterableDataset
import pandas as pd


class PandasIterableDataset(IterableDataset):
    def __init__(self, file_path):
        self.data_iter = pd.read_csv(file_path, iterator=True, header=None, chunksize=1)

    def __iter__(self):
        for data in self.data_iter:
            yield data


if __name__ == '__main__':
    dataset = PandasIterableDataset('test_csv.csv')
    for data in dataset:
        print(data)
Logo

为开发者提供学习成长、分享交流、生态实践、资源工具等服务,帮助开发者快速成长。

更多推荐