pytorch构造可迭代的Dataset——IterableDataset（pytorch Data学习二）

如果是可以一次性加载进内存的数据，上一篇博客：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)