当前位置： 首页 > article >正文

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

article 2025/3/9 9:12:07

使用 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 Image

class 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_dir
        self.transform = transform
        self.target_transform = target_transform

    def __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.csv
    img_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_dir
        self.transform = transform
        self.target_transform = target_transform

    def __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.0
    for 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 = 0
    total = 0
    val_loss = 0.0
    with 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 / total
    
    print(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 后对验证集进行评估可以及时发现过拟合和调参问题；最终对测试集进行评估可以获得模型的实际泛化性能。

http://www.kler.cn/a/471660.html

相关文章：

  • Mysql--基础篇--数据类型（整数，浮点数，日期，枚举，二进制，空间类型等）
  • Jupyter Markdown样式说明
  • HTML静态网页成品作业(HTML+CSS)——婚礼婚纱网页设计制作(6个页面)
  • OSPF使能配置
  • 攻防靶场(32)：两个爆破技巧 Funbox 7 EasyEnum
  • 5. 多线程(3) --- synchronized
  • 力扣经典题目之2283. 判断一个数的数字计数是否等于数位的值
  • [网络安全]DVWA之File Upload—AntSword(蚁剑)攻击姿势及解题详析合集
  • windows蓝屏以及windows补丁回滚
  • Solaris操作系统
  • JavaScript系列（12）-- 高阶函数应用
  • Spring5框架之SpringMVC
  • 记录一下Coding一直不能clone
  • 7_TypeScript Number --[深入浅出 TypeScript 测试]
  • Ae：合成设置 - 3D 渲染器
  • 除了RAII和智能指针，还有哪些资源管理机制？
  • Java设计模式 —— 【行为型模式】命令模式（Command Pattern） 详解
  • doris 2.1 Data Queries Common Table Expression UDF 学习笔记
  • 【LeetCode】4. 去重的效率提升
  • 基于CentOS的Docker + Nginx + Gitee + Jenkins部署总结（进阶）-- 接入钉钉通知功能