如何在PyTorch中加载个人数值类型CSV格式的自定义数据集

你已经熟悉内置数据集的使用逻辑，最适配的方案是基于torch.utils.data的标准接口做轻量改造，后续DataLoader的用法和内置数据集完全一致，没有额外学习成本。

前置依赖 #

需要提前安装pandas用于读取CSV：

pip install pandas

方案1：极简无自定义类方案（适合纯数值、无需复杂预处理的场景） #

不需要写任何类，几行代码就能完成加载：

import pandas as pd
import torch
from torch.utils.data import TensorDataset, DataLoader

# 读取CSV，拆分特征和标签
df = pd.read_csv("你的数据集路径.csv")
# 假设标签列名为label，drop之后剩下的都是特征列
features = torch.tensor(df.drop("label", axis=1).values.astype("float32"))
labels = torch.tensor(df["label"].values)

# 封装为PyTorch标准数据集，用法和内置数据集完全一致
dataset = TensorDataset(features, labels)
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

方案2：自定义Dataset类（更灵活，支持后续加预处理逻辑） #

如果后续需要加归一化、缺失值处理等逻辑，推荐用这个方案：

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

class CSVDataset(Dataset):
    def __init__(self, csv_path, label_col):
        # 读取CSV文件，无表头的话加header=None参数
        self.df = pd.read_csv(csv_path)
        # 拆分特征和标签，标签是索引的话用iloc取，比如最后一列：self.df.iloc[:, -1]
        self.labels = self.df[label_col].values
        self.features = self.df.drop(label_col, axis=1).values.astype("float32")

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

    def __getitem__(self, idx):
        # 返回单条样本的特征和标签张量
        return torch.tensor(self.features[idx]), torch.tensor(self.labels[idx])

使用方法和内置数据集完全一致：

# 初始化数据集
train_dataset = CSVDataset(csv_path="./train.csv", label_col="label")
# 拆分训练验证集的方法和内置数据集完全相同
train_set, val_set = torch.utils.data.random_split(train_dataset, [0.8, 0.2])
# 封装为DataLoader
train_loader = DataLoader(train_set, batch_size=32, shuffle=True)

# 迭代逻辑和内置数据集无差异
for batch_x, batch_y in train_loader:
    # 你的训练逻辑
    pass

注意事项 #

• 分类任务标签可以用torch.long类型，回归任务标签改成torch.float32即可
• 如果CSV没有表头，读取时加pd.read_csv(csv_path, header=None)，标签用列索引取值即可

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