如何高效将大量.pt字典文件中的Tensor加载到DataLoader？

方案1：按需读取单个.pt文件（推荐，低内存占用） #

不要一次性把所有20GB数据加载到内存，而是通过自定义Dataset类，在获取样本时才读取对应的.pt文件，完美适配DataLoader的流式处理逻辑，无需处理大字典的keys/values转列表问题。

代码实现：

import torch
from torch.utils.data import Dataset, DataLoader
import glob

class PTFileDataset(Dataset):
    def __init__(self, pt_dir):
        # 收集所有.pt文件路径（支持子目录递归查找）
        self.file_paths = glob.glob(f"{pt_dir}/**/*.pt", recursive=True)
    
    def __len__(self):
        return len(self.file_paths)
    
    def __getitem__(self, idx):
        # 加载单个.pt文件的字典
        sample_dict = torch.load(self.file_paths[idx])
        # 提取数据（字典的键）和标签（字典的值）
        data = next(iter(sample_dict.keys()))
        label = next(iter(sample_dict.values()))
        return data, label

# 初始化数据集和DataLoader
dataset = PTFileDataset("your/pt/files/directory")
dataloader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4,  # 根据CPU核心数调整，加速数据加载
    pin_memory=True  # 用GPU训练时开启，提升数据传输效率
)

# 使用示例
for batch_data, batch_labels in dataloader:
    # 执行训练/推理逻辑
    pass

方案2：基于已合并的大字典（内存足够时使用） #

如果你已经把所有字典合并成了一个大字典，直接将数据和标签分别堆叠成Tensor，用TensorDataset对接DataLoader：

import torch
from torch.utils.data import TensorDataset, DataLoader

# 假设all_data_dict是你合并后的大字典
data_tensor = torch.stack(list(all_data_dict.keys()))
label_tensor = torch.stack(list(all_data_dict.values()))

dataset = TensorDataset(data_tensor, label_tensor)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

关键优势说明 #

• 方案1的按需读取避免了一次性加载20GB数据导致的内存溢出问题，适合大体积数据集。
• 两种方案都不需要额外转换字典的keys/values为列表后再处理，直接完成数据到DataLoader的对接。

内容的提问来源于stack exchange，提问作者dhgGaurab