如何让自定义DataLoader分别返回特征值与对应标签？

原始代码错误点 #

• 不能直接将自定义Dataset实例传入torch.tensor()，DataLoader的第一个参数直接接收Dataset对象即可
• 原__getitem__返回的是 pandas Series 对象，没有拆分特征和标签，无法直接在遍历时解包为特征和标签
• 没有将返回的数值转为张量格式，无法匹配 PyTorch 训练的常规输入要求

修改后的完整代码 #

import torch
import pandas as pd
from torch.utils.data import Dataset
from torch.utils.data import DataLoader

class PandasDataset(Dataset):
    def __init__(self, dataframe, feature_cols=['values', 'values2'], label_col='labels'):
        self.features = dataframe[feature_cols].values  # 提前提取特征数组
        self.labels = dataframe[label_col].values       # 提前提取标签数组

    def __len__(self):
        # 返回数据集总样本数，DataLoader靠这个划分批次
        return len(self.features)

    def __getitem__(self, index):
        # 按索引取单条样本，返回特征张量+标签张量
        return torch.tensor(self.features[index], dtype=torch.float32), \
               torch.tensor(self.labels[index], dtype=torch.long)


d = {'values': [1, 2], 'values2': [3, 4],'labels': [5, 6]}
df = pd.DataFrame(data=d)
dataset = PandasDataset(df)
# 直接传入dataset，不需要转tensor
loader = DataLoader(dataset, batch_size=1, shuffle=False)

for batch_index, (values, label) in enumerate(loader):
    print(f"批次{batch_index}特征：", values)
    print(f"批次{batch_index}标签：", label)

运行原理说明 #

• Dataset是 PyTorch 规定的数据集抽象类，只要实现__len__和__getitem__两个魔法方法即可正常被DataLoader调用
  • __len__返回整个数据集的样本总数，DataLoader根据这个值和batch_size计算总批次数量
  • __getitem__接收索引值，返回对应位置的单条样本，DataLoader会自动把多个单条样本拼接为批次张量
• 我们在__init__里提前拆分特征列和标签列，避免每次取样本的时候重复做列筛选，运行效率更高
• 直接在__getitem__里把numpy数组转为PyTorch张量，不需要额外做类型转换即可直接传入模型训练
• 如果你的数据集特征列名/标签列名不同，只需要在初始化PandasDataset的时候传入对应的feature_cols和label_col参数即可适配

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