【深度学习】Pytorch：加载自定义数据集

本教程将使用 flower_photos 数据集演示如何在 PyTorch 中加载和导入自定义数据集。该数据集包含不同花种的图像，每种花的图像存储在以花名命名的子文件夹中。我们将深入讲解每个函数和对象的使用方法，使读者能够推广应用到其他数据集任务中。

flower_photos/
├── daisy/
│   ├── image1.jpg
│   ├── image2.jpg
└── rose/
     ├── image1.jpg
     ├── image2.jpg
...

环境配置

所需工具和库

pip install torch torchvision matplotlib

导入必要的库

import os
import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
from PIL import Image
import pathlib

数据集导入方法

定义数据转换

图像转换在计算机视觉任务中至关重要。通过 transforms 对象，我们可以实现图像大小调整、归一化、随机变换等预处理操作。

# 定义图像转换  
transform = transforms.Compose([  
    transforms.Resize((150, 150)),  # 调整图像大小为 150x150  
    transforms.ToTensor(),  # 将图像转换为 PyTorch 张量  
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])  # 归一化图像数据  
])  

# 数据路径  
data_dir = r"E:\CodeSpace\Deep\data\flower_photos"  

# 使用 ImageFolder 加载数据  
full_dataset = datasets.ImageFolder(root=data_dir, transform=transform)  

# 计算训练集和测试集的样本数量（80%和20%的划分）  
train_size = int(0.8 * 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=32, shuffle=True)  
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)  

# 获取类别名  
class_names = full_dataset.classes  
print("类别名:", class_names)

显示部分样本图像

可视化样本数据有助于理解数据集结构和数据质量。

# 定义函数以绘制样本图像
def plot_images(images, labels, class_names):
    plt.figure(figsize=(10, 10))
    for i in range(9):  # 绘制前 9 张图像
        plt.subplot(3, 3, i + 1)
        img = images[i].permute(1, 2, 0)  # 将张量维度从 (C, H, W) 转为 (H, W, C)
        plt.imshow(img * 0.5 + 0.5)  # 反归一化处理，恢复到原始像素范围 [0, 1]
        plt.title(class_names[labels[i]])  # 显示类别标签
        plt.axis('off')  # 去掉坐标轴

# 获取部分样本数据用于展示
sample_images, sample_labels = next(iter(train_loader))
plot_images(sample_images, sample_labels, class_names)

自定义数据加载方法

当数据结构复杂或需要额外处理时，可以通过继承 torch.utils.data.Dataset 创建自定义数据加载类。

Dataset 类详解

Dataset 是 PyTorch 中的一个抽象类，用户需要实现以下核心方法：

1. __init__()：初始化方法
   • 传入数据路径和转换方法。
   • 加载所有图像路径并生成类别标签。
2. __len__()：返回数据集大小
   • 指定数据集中样本数量。
3. __getitem__()：根据索引获取样本数据
   • 加载指定位置的图像和标签，并进行必要的转换。

代码实现

class CustomFlowerDataset(torch.utils.data.Dataset):
    def __init__(self, data_dir, transform=None):
        # 初始化数据集路径和图像转换方法
        self.data_dir = pathlib.Path(data_dir)
        self.transform = transform
        self.image_paths = list(self.data_dir.glob('*/*.jpg'))  # 获取所有图像路径
        self.label_names = sorted(item.name for item in self.data_dir.glob('*/') if item.is_dir())
        self.label_to_index = {name: idx for idx, name in enumerate(self.label_names)}  # 将类别名映射为索引

    def __len__(self):
        # 返回数据集大小
        return len(self.image_paths)

    def __getitem__(self, idx):
        # 根据索引获取图像及其标签
        img_path = self.image_paths[idx]
        label = self.label_to_index[img_path.parent.name]  # 通过父文件夹名获取标签
        image = Image.open(img_path).convert("RGB")  # 确保图像是 RGB 模式
        if self.transform:
            image = self.transform(image)  # 进行图像预处理
        return image, label

# 使用自定义数据集
custom_dataset = CustomFlowerDataset(data_dir, transform=transform)
custom_loader = DataLoader(custom_dataset, batch_size=32, shuffle=True)

随机划分数据集

如果你还希望在这个自定义数据集上随机划分训练集和测试集，可以使用 torch.utils.data.random_split。以下是示例代码：

from torch.utils.data import random_split  

# 获取数据集长度  
full_dataset = CustomFlowerDataset(data_dir, transform=transform)  

# 计算训练集和测试集的样本数量（80%和20%的划分）  
train_size = int(0.8 * 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=32, shuffle=True)  
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)  

print(f"训练集大小: {len(train_dataset)}, 测试集大小: {len(test_dataset)}")  

数据加载性能优化

• num_workers 参数：设置并行数据加载线程数。对于多核 CPU，可以显著提高数据加载效率。

• prefetch_factor 参数：控制每个工作线程预取的批次数量。

  custom_loader = DataLoader(custom_dataset, batch_size=32, shuffle=True, num_workers=4, prefetch_factor=2)

Dataset 类扩展建议

1. 支持多格式数据读取 ：通过扩展 __getitem__() 来支持其他格式如 PNG、BMP。
2. 数据过滤 ：在 __init__() 中根据文件名或元数据筛选特定样本。
3. 标签增强：为每个样本生成附加信息，例如图像的元数据或分布特征。

数据集的使用方法

遍历数据集

模型训练前需要遍历数据集以加载图像和标签：

for images, labels in custom_loader:
    # images 是图像张量，labels 是对应的类别标签
    print(f"图像张量大小: {images.shape}, 标签: {labels}")

模型输入

数据集加载完成后可直接用于模型训练：

import torch.nn as nn
import torch.optim as optim

# 定义一个简单的神经网络模型
model = nn.Sequential(
    nn.Flatten(),  # 将输入张量展平成一维
    nn.Linear(150*150*3, 128),  # 输入层到隐藏层的全连接层
    nn.ReLU(),  # 激活函数
    nn.Linear(128, len(class_names))  # 输出层，类别数量等于花的种类数
)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()  # 交叉熵损失适用于多分类问题
optimizer = optim.Adam(model.parameters(), lr=0.001)  # Adam 优化器

# 示例训练过程
for epoch in range(2):  # 简单训练两轮
    for images, labels in custom_loader:
        outputs = model(images)  # 前向传播计算输出
        loss = criterion(outputs, labels)  # 计算损失

        optimizer.zero_grad()  # 梯度清零
        loss.backward()  # 反向传播计算梯度
        optimizer.step()  # 更新模型参数

    print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")

模型评估

加载后的数据集也可用于验证模型性能：

correct = 0
total = 0
model.eval()  # 设置模型为评估模式
with torch.no_grad():
    for images, labels in test_loader:
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

accuracy = 100 * correct / total
print(f"模型准确率: {accuracy:.2f}%")

方法对比与扩展

ImageFolder vs 自定义 Dataset

• ImageFolder：适合简单目录结构，快速加载标准图像数据。
• 自定义 Dataset：更适合复杂数据结构及自定义逻辑，例如多模态数据处理。

提高模型泛化能力

• 数据增强 ：通过 transforms.RandomHorizontalFlip()、transforms.ColorJitter() 等方法增加数据多样性。
• 归一化技巧 ：根据数据集的特性调整 mean 和 std 参数。

总结

本教程详细讲解了如何在 PyTorch 中加载和导入 flower_photos 数据集，结合不同方法的讲解使你能根据项目需求灵活选择适合的数据加载方案。同时，我们探讨了优化和扩展方法，希望这些内容能为你的深度学习项目提供有力支持。