如何在PyTorch中使用带prefetch_factor的DataLoader加载超大txt数据集？

1. 构建支持逐行读取的自定义Dataset #

因为数据无法全量载入内存，所以需要实现按偏移读取的Dataset，避免加载全量数据。首先需要提前扫描txt文件生成每行的偏移索引，后续可以直接通过偏移量快速定位到对应行，不用每次遍历全文件。

import os
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader

# 生成行偏移索引，仅需执行一次，后续可以直接加载生成好的索引文件
def build_line_index(txt_path, index_save_path="line_offset_index.npy"):
    if os.path.exists(index_save_path):
        return np.load(index_save_path)
    line_offsets = [0]
    with open(txt_path, 'rb') as f:
        while f.readline():
            line_offsets.append(f.tell())
    # 移除最后一个文件末尾的偏移
    line_offsets = line_offsets[:-1]
    np.save(index_save_path, line_offsets)
    return line_offsets

# 自定义数据集，按索引读取对应行数据
class LargeTxtDataset(Dataset):
    def __init__(self, txt_path, line_index, delimiter=' ', transform=None):
        self.txt_path = txt_path
        self.line_index = line_index
        self.delimiter = delimiter
        self.transform = transform
        # 每个worker初始化时会单独持有一个文件句柄，不存在多进程冲突
        self.file_handler = open(txt_path, 'rb')

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

    def __getitem__(self, idx):
        # 直接跳转到对应行的偏移位置读取
        self.file_handler.seek(self.line_index[idx])
        line = self.file_handler.readline().decode('utf-8').strip()
        # 解析为浮点数组
        data = np.array(line.split(self.delimiter), dtype=np.float32)
        if self.transform:
            data = self.transform(data)
        return torch.from_numpy(data)

2. 配置带prefetch_factor的DataLoader #

prefetch_factor参数仅在num_workers > 0时生效，含义是每个工作进程预取到内存的批次数量，数值越高预存的批次越多，训练时IO等待时间越少，但内存占用也越高。

# 加载行索引
line_index = build_line_index("你的大文件路径.txt")
dataset = LargeTxtDataset("你的大文件路径.txt", line_index, delimiter=' ')

# 初始化带预取功能的DataLoader
dataloader = DataLoader(
    dataset,
    batch_size=256, # 根据你的显存大小调整
    num_workers=4, # 必须大于0，prefetch_factor才生效，机械硬盘建议2~4，固态硬盘建议4~8
    prefetch_factor=8, # 每个worker预取8个批次，内存够用可以调到16，默认值为2
    shuffle=True,
    pin_memory=True # 用GPU训练时开启，加快数据转移速度
)

注意事项 #

• 行索引只需生成一次，后续训练直接加载npy格式的索引文件即可，不用重复扫描全量txt
• 纯txt的IO效率很低，如果允许的话建议提前将txt转成HDF5或分块npy格式存储，读取速度能提升5~10倍
• num_workers不要设置过高，否则多进程同时读盘会引发IO竞争反而降低速度
• 如果需要加载标签，可以在Dataset中额外加入标签读取逻辑，和数据读取逻辑保持一致即可

内容的提问来源于stack exchange，提问作者G-09