本文将为您介绍如何在 PyTorch 中进行数据加载，包括常见的数据集加载方法以及如何自定义数据加载器。

常见数据集加载

PyTorch 提供了多种数据集加载方法，以下是一些常用的数据集：

  • CIFAR-10:

    • CIFAR-10
    • CIFAR-10 是一个包含 10 个类别的 32x32 彩色图像数据集，共有 60000 张训练图像和 10000 张测试图像。

  • MNIST:

    • MNIST
    • MNIST 是一个手写数字数据集，包含 60,000 个训练样本和 10,000 个测试样本。

  • ImageFolder:

    • ImageFolder
    • ImageFolder 是一个文件夹，其中包含图像文件，PyTorch 可以自动将其加载为数据集。

自定义数据加载器

如果您需要加载特定格式的数据集，可以自定义数据加载器。以下是一个简单的自定义数据加载器示例：

from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import os

class CustomDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.image_paths = [os.path.join(root_dir, f) for f in os.listdir(root_dir)]

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

    def __getitem__(self, idx):
        img_path = self.image_paths[idx]
        image = Image.open(img_path)
        if self.transform:
            image = self.transform(image)
        return image

# 使用自定义数据加载器
transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
])

dataset = CustomDataset(root_dir='path_to_your_dataset', transform=transform)
dataloader = DataLoader(dataset, batch_size=4, shuffle=True)

扩展阅读