快速读取(小数据)
先用 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()。
