PyTorch 自定义数据集示例教程

自定义数据集是机器学习中一个非常重要的概念，尤其是在使用深度学习框架如PyTorch时。本文将为您介绍如何使用PyTorch创建自定义数据集。

创建自定义数据集

在PyTorch中，您可以通过继承torch.utils.data.Dataset类来创建自定义数据集。以下是一个简单的示例：

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

class CustomDataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.data = os.listdir(root_dir)

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

    def __getitem__(self, idx):
        img_name = os.path.join(self.root_dir, self.data[idx])
        image = Image.open(img_name)
        if self.transform:
            image = self.transform(image)
        return image

在上面的代码中，我们创建了一个名为CustomDataset的自定义数据集类。它接受一个根目录路径和一个可选的转换函数。__len__方法返回数据集中的样本数量，而__getitem__方法用于获取索引为idx的样本。

使用数据加载器

创建自定义数据集后，您可以使用DataLoader来加载数据。DataLoader可以自动批处理数据，并且可以并行加载数据。

transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
])

dataset = CustomDataset(root_dir='path/to/your/data', transform=transform)
dataloader = DataLoader(dataset, batch_size=4, shuffle=True)

for images, labels in dataloader:
    # 在这里进行模型训练或其他操作
    pass

在上面的代码中，我们首先创建了一个转换函数，它将图像调整到256x256的尺寸，并将其转换为PyTorch张量。然后，我们创建了一个CustomDataset实例，并使用DataLoader来加载数据。

扩展阅读

如果您想了解更多关于PyTorch的内容，请访问我们的PyTorch教程页面。

Image