如何使用PyTorch的random_split()拆分数据集并喂入模型？

核心问题梳理 #

1. 未分离特征与标签：原代码直接对Pandas DataFrame进行拆分，没有将输入特征和预测标签分开，不符合PyTorch模型的输入要求。
2. Dataset/DataLoader使用错误：random_split不能直接作用于DataFrame，需基于PyTorch标准Dataset类处理数据；且不能将DataLoader直接传入模型forward方法，模型需要的是张量输入，DataLoader是批量迭代数据的容器。
3. 模型forward无返回值：原模型的forward方法没有输出结果，无法计算损失和完成训练流程。

修正后的完整代码 #

1. 重构数据集类（dataset.py） #

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

class PimaDataset(Dataset):
    divide_rate = 0.8
    file = './pima-indians-diabetes.csv'

    def __init__(self):
        # 读取数据并分离特征与标签（假设最后一列为标签）
        data_set = pd.read_csv(self.file)
        # 特征转为浮点型张量
        self.features = torch.tensor(data_set.iloc[:, :-1].values, dtype=torch.float32)
        # 标签转为二维浮点型张量（匹配模型输出维度）
        self.labels = torch.tensor(data_set.iloc[:, -1].values, dtype=torch.float32).unsqueeze(1)

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

    def __getitem__(self, idx):
        # 返回单条数据的特征与标签
        return self.features[idx], self.labels[idx]

def get_dataloaders(batch_size=32):
    # 加载完整数据集
    full_dataset = PimaDataset()
    # 拆分训练集与测试集
    train_size = int(PimaDataset.divide_rate * len(full_dataset))
    test_size = len(full_dataset) - train_size
    train_dataset, test_dataset = random_split(full_dataset, [train_size, test_size])
    
    # 创建数据加载器，负责批量迭代与洗牌
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
    
    return train_loader, test_loader

2. 修正模型与训练逻辑（主文件） #

import torch
import torch.nn as nn
from torch.optim import SGD
from dataset import get_dataloaders

class NeuralNetwork(nn.Module):
    def __init__(self, input_dim=8, hidden_dim=4, output_dim=1):
        super().__init__()
        # 用Sequential简化层结构
        self.layers = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.RReLU(),
            nn.Linear(hidden_dim, output_dim)
        )

    def forward(self, x):
        # 返回模型输出结果
        return self.layers(x)

# 获取训练/测试数据加载器
train_loader, test_loader = get_dataloaders()

# 初始化模型、优化器、损失函数（二分类任务用BCEWithLogitsLoss）
model = NeuralNetwork()
optimizer = SGD(model.parameters(), lr=0.01)
criterion = nn.BCEWithLogitsLoss()

# 训练循环示例
epochs = 10
model.train()
for epoch in range(epochs):
    total_loss = 0.0
    # 遍历DataLoader获取批量特征与标签
    for features, labels in train_loader:
        optimizer.zero_grad()
        # 传入特征张量，得到模型输出
        outputs = model(features)
        # 计算损失
        loss = criterion(outputs, labels)
        # 反向传播+参数更新
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    print(f"Epoch {epoch+1}/{epochs}, 平均损失: {total_loss/len(train_loader):.4f}")

# 测试模型准确率
model.eval()
correct = 0
total = 0
with torch.no_grad():
    for features, labels in test_loader:
        outputs = model(features)
        # 将logits转为概率后取预测结果
        predicted = torch.round(torch.sigmoid(outputs))
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
print(f"测试集准确率: {100 * correct / total:.2f}%")

关键修正说明 #

• 特征与标签分离：在自定义Dataset类中，明确提取输入特征（除最后一列外的所有数据）和标签（最后一列），并转为PyTorch张量。
• 标准Dataset/DataLoader流程：基于torch.utils.data.Dataset构建数据集，再用random_split拆分，最后用DataLoader实现批量数据迭代。
• 模型输入输出修正：模型forward方法返回输出结果，且根据Pima数据集实际特征数（8个）调整输入维度，匹配任务需求。
• 训练逻辑规范：遍历DataLoader获取批量数据，完成前向传播、损失计算、反向传播和参数更新的完整训练流程。

内容的提问来源于stack exchange，提问作者M a m a D