如何在PyTorch中定义data loader？已熟悉TensorFlow对应写法

你给出的TensorFlow示例是直接从内存张量构造数据集，PyTorch中可以通过TensorDataset+DataLoader的组合实现完全对等的功能，代码如下：

首先导入依赖：

from torch.utils.data import TensorDataset, DataLoader

构造训练、验证加载器：

# 训练集加载器
train_dataset = TensorDataset(x_train, y_train)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

# 验证集加载器
validation_dataset = TensorDataset(x_val, y_val)
validation_loader = DataLoader(validation_dataset, batch_size=32, shuffle=False)

参数说明 #

DataLoader支持配置常用的加载参数：

• batch_size：每批次返回的样本数量，可根据显存大小调整
• shuffle：是否打乱样本顺序，训练集建议开启，验证/测试集关闭
• num_workers：加载数据的子进程数，Windows环境出现启动错误时可设为0
• drop_last：是否丢弃最后一批不足batch_size的样本

自定义数据集场景 #

如果你的数据不是预加载到内存的张量，需要从本地文件读取、或者需要自定义预处理逻辑，可以继承Dataset基类实现自定义数据集：

from torch.utils.data import Dataset

class MyCustomDataset(Dataset):
    def __init__(self, sample_list, label_list):
        # 初始化方法，可传入文件路径列表、预处理规则等
        self.samples = sample_list
        self.labels = label_list

    def __len__(self):
        # 返回数据集总样本数
        return len(self.samples)

    def __getitem__(self, idx):
        # 按索引返回单条样本，可在此处实现读文件、数据增强等逻辑
        return self.samples[idx], self.labels[idx]

# 自定义数据集构造加载器的方法和内置数据集一致
train_custom_dataset = MyCustomDataset(x_train, y_train)
train_custom_loader = DataLoader(train_custom_dataset, batch_size=32, shuffle=True)

使用方法 #

直接迭代加载器即可按批次获取数据：

for batch_x, batch_y in train_loader:
    # 此处写入模型训练/验证逻辑，batch_x、batch_y为当前批次的张量
    pass

内容的提问来源于stack exchange，提问作者Anox Software technology