PyTorch 中如何自定义 Dataset？

推荐答案

在 PyTorch 中，自定义 Dataset 需要继承 torch.utils.data.Dataset 类，并实现 __len__ 和 __getitem__ 方法。以下是一个简单的示例：

// www.javascriptcn.com code example
import torch
from torch.utils.data import Dataset

class CustomDataset(Dataset):
    def __init__(self, data, labels):
        self.data = data
        self.labels = labels

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

    def __getitem__(self, idx):
        sample = self.data[idx]
        label = self.labels[idx]
        return sample, label

# 示例数据
data = torch.randn(100, 3, 32, 32)  # 100个3x32x32的随机张量
labels = torch.randint(0, 10, (100,))  # 100个随机标签

# 创建自定义Dataset实例
custom_dataset = CustomDataset(data, labels)

# 使用DataLoader加载数据
from torch.utils.data import DataLoader
dataloader = DataLoader(custom_dataset, batch_size=4, shuffle=True)

# 遍历DataLoader
for batch in dataloader:
    samples, labels = batch
    print(samples.shape, labels.shape)

本题详细解读

1. 继承 torch.utils.data.Dataset 类

在 PyTorch 中，Dataset 是一个抽象类，表示数据集的抽象。自定义数据集类需要继承 Dataset 类，并实现 __len__ 和 __getitem__ 方法。

2. 实现 __len__ 方法

__len__ 方法返回数据集的大小，即数据集中样本的数量。这个方法在调用 len(dataset) 时会被调用。

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

3. 实现 __getitem__ 方法

__getitem__ 方法根据索引 idx 返回数据集中的一个样本。这个方法在通过索引访问数据集时会被调用，例如 dataset[idx]。

def __getitem__(self, idx):
    sample = self.data[idx]
    label = self.labels[idx]
    return sample, label

4. 使用 DataLoader 加载数据

DataLoader 是 PyTorch 提供的一个工具，用于将数据集包装成一个可迭代的对象，支持批量加载、打乱数据、多线程加载等功能。

from torch.utils.data import DataLoader
dataloader = DataLoader(custom_dataset, batch_size=4, shuffle=True)

5. 遍历 DataLoader

通过 DataLoader，可以方便地遍历数据集中的每一个批次。每个批次包含多个样本和对应的标签。

for batch in dataloader:
    samples, labels = batch
    print(samples.shape, labels.shape)

6. 自定义数据集的灵活性

自定义 Dataset 类的灵活性在于可以根据实际需求调整 __getitem__ 方法。例如，可以从文件中加载数据、进行数据预处理、数据增强等操作。

def __getitem__(self, idx):
    sample = self.data[idx]
    label = self.labels[idx]
    # 数据预处理
    sample = self.transform(sample)
    return sample, label

通过这种方式，可以轻松地将自定义数据集集成到 PyTorch 的训练流程中。