PyTorch 读取 CSV 数据:pandas 转 Tensor、自定义 Dataset 与 DataLoader

快速读取(小数据)

先用 pandas 读 CSV,再一次性转成 Tensor:

import pandas as pd
import torch

df = pd.read_csv("house.csv")

# 前 N-1 列是特征,最后一列是标签
x = torch.tensor(df.iloc[:, :-1].values, dtype=torch.float32)
y = torch.tensor(df.iloc[:, -1].values, dtype=torch.float32)

print(x.shape, y.shape)  # torch.Size([N, features]) torch.Size([N])

iloc[:, :-1] 取除最后列之外的所有列,iloc[:, -1] 取最后一列。

自定义 Dataset(正式训练推荐)

当数据量大或需要在线增强时,继承 Dataset 类实现懒加载:

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

class CSVDataset(Dataset):
    def __init__(self, path):
        self.data = pd.read_csv(path)

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

    def __getitem__(self, index):
        row = self.data.iloc[index]
        x = torch.tensor(row[:-1].values, dtype=torch.float32)
        y = torch.tensor(row[-1], dtype=torch.float32)
        return x, y

必须实现三个方法:

  • __init__:加载/读取数据
  • __len__:返回样本总数
  • __getitem__:返回第 index 个样本

DataLoader 分批读取

from torch.utils.data import DataLoader

dataset = CSVDataset("house.csv")

loader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4    # 多进程预取,Windows 下注意加 if __name__ == "__main__"
)

for x_batch, y_batch in loader:
    print(x_batch.shape)   # [32, features]
    print(y_batch.shape)   # [32]

DataLoader 自动处理:

  • 打乱数据(shuffle=True
  • 按 batch_size 分批
  • 多线程预取(num_workers

完整训练示例

import torch
import torch.nn as nn
from torch.utils.data import DataLoader

model = nn.Sequential(
    nn.Linear(2, 16),
    nn.ReLU(),
    nn.Linear(16, 1)
)

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.MSELoss()

dataset = CSVDataset("house.csv")
loader = DataLoader(dataset, batch_size=8, shuffle=True)

for epoch in range(100):
    for x, y in loader:
        pred = model(x).squeeze()
        loss = criterion(pred, y)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    if epoch % 10 == 0:
        print(f"Epoch {epoch}, Loss: {loss.item():.4f}")

训练集与验证集分割

from torch.utils.data import random_split

dataset = CSVDataset("house.csv")

train_size = int(0.8 * len(dataset))
val_size = len(dataset) - train_size

train_set, val_set = random_split(dataset, [train_size, val_size])

train_loader = DataLoader(train_set, batch_size=32, shuffle=True)
val_loader = DataLoader(val_set, batch_size=32)

处理标签类型

回归任务用 float32,分类任务用 long

# 回归:y 是连续值
y = torch.tensor(row[-1], dtype=torch.float32)

# 分类:y 是类别编号
y = torch.tensor(int(row[-1]), dtype=torch.long)

分类时对应的 loss 是 nn.CrossEntropyLoss(),回归时用 nn.MSELoss()