使用 PyTorch 自定义数据集并划分训练、验证与测试集

使用 PyTorch 自定义数据集并划分训练、验证与测试集

在图像分类等任务中，通常需要将原始训练数据进一步划分为训练集和验证集，以便在训练过程中评估模型的性能。下面将详细介绍如何组织数据与注释文件、如何分割训练集和验证集，以及如何基于自定义 Dataset 类构建 DataLoader 以加速模型训练与评估。

一、数据准备

1.1 文件结构

假设你的数据目录结构如下所示：

data/
├── train_data/
│   ├── img1.png
│   ├── img2.png
│   ├── img3.png
│   └── ...
├── test_data/
│   ├── img101.png
│   ├── img102.png
│   ├── img103.png
│   └── ...
├── train_annotations.csv
└── test_annotations.csv

注意：这里将 train_annotations.csv 和 test_annotations.csv 文件单独放在 data/ 目录下，而不放在各自图片的子文件夹中。这样当图片数量非常多时，我们也能快速找到并管理这两个 CSV 文件。

1.2 注释文件（CSV）格式示例

在 train_annotations.csv 和 test_annotations.csv 中，一般会包含两列或更多列信息，但最关键的通常是 图片文件名（filename）和 标签（label）。格式示例如下：

train_annotations.csv：

filename,label
img1.png,0
img2.png,1
img3.png,0
...

test_annotations.csv：

filename,label
img101.png,0
img102.png,1
img103.png,0
...

• filename 列表示图像的文件名，需要与 train_data/ 或 test_data/ 文件夹下的文件一一对应。
• label 列表示图像所对应的类别或标签，可以是整数，也可以是字符串，比如 cat、dog 等。训练时通常会将字符串映射到整数标签或独热编码。

二、将训练数据划分为训练集和验证集

在进行模型训练前，往往需要将原始训练数据（以下简称 “总训练集”）拆分成 训练集（train） 和 验证集（val）。这里我们使用 scikit-learn 提供的 train_test_split 函数来完成这一步骤。

import pandas as pd
from sklearn.model_selection import train_test_split# 读取原始训练集的注释文件（此时还未拆分）
train_annotations = pd.read_csv('data/train_annotations.csv')# 按 80%:20% 的比例拆分为 新的训练集(train_df) 和 验证集(val_df)
train_df, val_df = train_test_split(train_annotations, test_size=0.2, random_state=42, stratify=train_annotations['label']
)# 将拆分后的注释文件保存为新的 CSV 文件
train_df.to_csv('data/train_split.csv', index=False)
val_df.to_csv('data/val_split.csv', index=False)

关键参数说明：

• test_size=0.2：表示将 20% 的样本作为验证集，其余 80% 作为新的训练集。
• random_state=42：让划分结果可复现，方便后续对比不同实验结果。
• stratify=train_annotations['label']：在划分时保持各类别在训练和验证集中相同比例，这在分类任务中尤为重要。

执行完以上步骤后，你的 data 目录下会多出两个新的注释文件：

data/
├── train_data/
│   ├── ...
├── test_data/
│   ├── ...
├── train_annotations.csv   # 原始，总训练集注释
├── train_split.csv         # 新的，训练集注释
└── val_split.csv           # 新的，验证集注释

三、自定义 Dataset 类

PyTorch 提供了 torch.utils.data.Dataset 作为数据集的抽象基类。我们可以通过继承并重写其中的方法，来实现灵活的数据加载逻辑。

下面的 CustomImageDataset 类支持通过 CSV 文件（包括你在上一步生成的 train_split.csv, val_split.csv 等）来读取图像与标签，并在取样本时进行必要的预处理操作。

import os
import pandas as pd
from torch.utils.data import Dataset
from PIL import Imageclass CustomImageDataset(Dataset):def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):"""初始化数据集。参数:annotations_file (str): CSV 文件路径，包含 (filename, label) 等信息img_dir (str): 存放图像的文件夹路径transform (callable, optional): 对图像进行转换和增强的函数或 transforms 组合target_transform (callable, optional): 对标签进行转换的函数"""self.img_labels = pd.read_csv(annotations_file)self.img_dir = img_dirself.transform = transformself.target_transform = target_transformdef __len__(self):"""返回整个数据集的样本数量。"""return len(self.img_labels)def __getitem__(self, idx):"""根据索引 idx 获取单个样本。返回:(image, label) 其中 image 可以是一个 PIL 图像或 Tensor，label 可以是整数或字符串"""# 1. 获取图像文件名与对应的标签img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])label = self.img_labels.iloc[idx, 1]# 2. 读取图像并转换为 RGB 模式（如果是灰度则可用 'L'）image = Image.open(img_path).convert('RGB')# 3. 对图像和标签进行必要的变换if self.transform:image = self.transform(image)if self.target_transform:label = self.target_transform(label)return image, label

四、创建训练集、验证集、测试集对应的 DataLoader

有了自定义 Dataset 后，就可以利用 PyTorch 自带的 DataLoader 来进行批量数据加载、随机打乱以及多线程读取数据等工作。以下示例展示了如何分别实例化 训练集、验证集 和 测试集 的 Dataset 对象，并为每个对象创建 DataLoader。

from torchvision import transforms
from torch.utils.data import DataLoader# 定义训练、验证/测试时所需的数据变换
train_transform = transforms.Compose([transforms.Resize((224, 224)),transforms.RandomHorizontalFlip(),  # 数据增强transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])
])val_test_transform = transforms.Compose([transforms.Resize((224, 224)),transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])
])# 实例化训练集 (train_dataset)
train_dataset = CustomImageDataset(annotations_file='data/train_split.csv',  # 注意这里不再是 data/train_annotations.csvimg_dir='data/train_data',transform=train_transform
)# 实例化验证集 (val_dataset)
val_dataset = CustomImageDataset(annotations_file='data/val_split.csv',img_dir='data/train_data',transform=val_test_transform
)# 实例化测试集 (test_dataset)
test_dataset = CustomImageDataset(annotations_file='data/test_annotations.csv',img_dir='data/test_data',transform=val_test_transform
)# 构建 DataLoader
train_loader = DataLoader(train_dataset,batch_size=64,shuffle=True,      # 训练时常使用 shuffle=True 来打乱顺序num_workers=4,     # 根据 CPU 核心数进行调整drop_last=True     # 避免最后一个 batch 样本数不足时带来的问题
)val_loader = DataLoader(val_dataset,batch_size=64,shuffle=False,num_workers=4,drop_last=False
)test_loader = DataLoader(test_dataset,batch_size=64,shuffle=False,num_workers=4,drop_last=False
)

通过使用 DataLoader，你就可以在训练和验证过程中以 批（batch）为单位获取数据，从而显著提升训练速度，并方便进行数据增强、随机打乱等操作。

五、完整示例脚本

下面给出一个相对完整的示例脚本，整合了数据拆分、自定义数据集加载以及构建 DataLoader 的主要流程。如果你愿意，可以将这些步骤拆分到不同的 Python 文件中，以保持项目结构清晰。

import os
import pandas as pd
from sklearn.model_selection import train_test_split
from torch.utils.data import Dataset, DataLoader
from PIL import Image
from torchvision import transforms# ========== 1. 数据集拆分函数 ========== #
def split_train_val(annotations_file, output_train_file, output_val_file, test_size=0.2, random_state=42):df = pd.read_csv(annotations_file)train_df, val_df = train_test_split(df, test_size=test_size, random_state=random_state, stratify=df['label'])train_df.to_csv(output_train_file, index=False)val_df.to_csv(output_val_file, index=False)# ========== 2. 定义自定义 Dataset 类 ========== #
class CustomImageDataset(Dataset):def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):self.img_labels = pd.read_csv(annotations_file)self.img_dir = img_dirself.transform = transformself.target_transform = target_transformdef __len__(self):return len(self.img_labels)def __getitem__(self, idx):img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])label = self.img_labels.iloc[idx, 1]image = Image.open(img_path).convert('RGB')if self.transform:image = self.transform(image)if self.target_transform:label = self.target_transform(label)return image, label# ========== 3. 执行划分并创建训练/验证/测试集 ========== #
# 假设原始的训练集标注文件位于 data/train_annotations.csv
split_train_val(annotations_file='data/train_annotations.csv',output_train_file='data/train_split.csv',output_val_file='data/val_split.csv',test_size=0.2,random_state=42
)train_transform = transforms.Compose([transforms.Resize((224, 224)),transforms.RandomHorizontalFlip(),transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])
])val_test_transform = transforms.Compose([transforms.Resize((224, 224)),transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])
])train_dataset = CustomImageDataset(annotations_file='data/train_split.csv',img_dir='data/train_data',transform=train_transform
)val_dataset = CustomImageDataset(annotations_file='data/val_split.csv',img_dir='data/train_data',transform=val_test_transform
)test_dataset = CustomImageDataset(annotations_file='data/test_annotations.csv',img_dir='data/test_data',transform=val_test_transform
)train_loader = DataLoader(train_dataset,batch_size=64,shuffle=True,num_workers=4,drop_last=True
)val_loader = DataLoader(val_dataset,batch_size=64,shuffle=False,num_workers=4,drop_last=False
)test_loader = DataLoader(test_dataset,batch_size=64,shuffle=False,num_workers=4,drop_last=False
)# ========== 4. 简单测试：读取一个 batch ========== #
for images, labels in train_loader:print(images.shape, labels.shape)break

六、在训练循环中使用验证集

构建好训练、验证和测试集的 DataLoader 之后，你就可以在模型训练过程中使用验证集来评估模型性能；并在完全训练结束后，对测试集进行最终评估。以下是一个最简化的示例，演示如何在每个 epoch 后进行验证：

import torch
import torch.nn as nn
import torch.optim as optim# 定义简单的神经网络
class SimpleNN(nn.Module):def __init__(self, num_classes=10):super(SimpleNN, self).__init__()self.flatten = nn.Flatten()self.fc1 = nn.Linear(224*224*3, 128)  # 根据输入图像大小进行调整self.relu = nn.ReLU()self.fc2 = nn.Linear(128, num_classes)def forward(self, x):x = self.flatten(x)x = self.relu(self.fc1(x))x = self.fc2(x)return x# 初始化模型、损失函数和优化器
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = SimpleNN(num_classes=2).to(device)  # 假设有 2 个类别
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)# 训练过程
num_epochs = 5
for epoch in range(num_epochs):# 1. 训练阶段model.train()running_loss = 0.0for images, labels in train_loader:images = images.to(device)labels = labels.to(device)optimizer.zero_grad()outputs = model(images)loss = criterion(outputs, labels)loss.backward()optimizer.step()running_loss += loss.item()avg_train_loss = running_loss / len(train_loader)# 2. 验证阶段model.eval()correct = 0total = 0val_loss = 0.0with torch.no_grad():for images, labels in val_loader:images = images.to(device)labels = labels.to(device)outputs = model(images)loss = criterion(outputs, labels)val_loss += loss.item()_, predicted = torch.max(outputs.data, 1)total += labels.size(0)correct += (predicted == labels).sum().item()avg_val_loss = val_loss / len(val_loader)val_accuracy = 100.0 * correct / totalprint(f'Epoch [{epoch+1}/{num_epochs}], 'f'Train Loss: {avg_train_loss:.4f}, 'f'Val Loss: {avg_val_loss:.4f}, 'f'Val Accuracy: {val_accuracy:.2f}%')

输出示例：

Epoch [1/5], Train Loss: 1.2034, Val Loss: 0.4567, Val Accuracy: 85.32%
Epoch [2/5], Train Loss: 0.9876, Val Loss: 0.3987, Val Accuracy: 88.45%
...

---

总结

1. 数据组织：将大量图片与注释文件分开存储（如 train_annotations.csv 与 test_annotations.csv 单独放在 data/ 目录下），可以在图片数量庞大时更方便地管理和检索。
2. 数据集拆分：使用 train_test_split 将原始训练集拆分为训练集与验证集，以便在训练过程中监控模型的过拟合情况。
3. 自定义 Dataset：通过继承 Dataset 并重写 __getitem__、__len__，可以灵活处理任意格式的数据，并在读入时执行预处理/增强操作。
4. 构建 DataLoader：使用 PyTorch 的 DataLoader 可以轻松实现批量读取、并行加速、随机打乱等功能，大幅提升训练效率。
5. 验证与测试：在每个 epoch 后对验证集进行评估可以及时发现过拟合和调参问题；最终对测试集进行评估可以获得模型的实际泛化性能。