pytorch中的dataset用法详解

在PyTorch中，torch.utils.data.Dataset是一个抽象类，用于表示数据集。我们可以使用torch.utils.data.Dataset类来加载和处理数据集。以下是两个示例说明。

示例1：自定义数据集

import torch
from torch.utils.data import Dataset

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

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

    def __getitem__(self, index):
        x = self.data[index][0]
        y = self.data[index][1]
        return x, y

# 定义数据集
data = [(1, 2), (3, 4), (5, 6), (7, 8)]
dataset = CustomDataset(data)

# 加载数据集
dataloader = torch.utils.data.DataLoader(dataset, batch_size=2, shuffle=True)

# 输出数据集
for batch in dataloader:
    x, y = batch
    print(x, y)

在这个示例中，我们首先定义了一个名为CustomDataset的自定义数据集类，该类继承自torch.utils.data.Dataset类。然后，我们在__init__函数中初始化数据集，并在__len__函数中返回数据集的长度。最后，我们在__getitem__函数中返回数据集中的一个样本。

接下来，我们定义了一个名为data的数据集，并使用CustomDataset类将其转换为数据集对象。然后，我们使用torch.utils.data.DataLoader函数加载数据集，并使用for循环遍历数据集中的每个batch，并输出每个batch中的数据。

示例2：使用现有数据集

import torch
import torchvision
import torchvision.transforms as transforms

# 定义transforms
transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

# 加载数据集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
                                          shuffle=True, num_workers=2)

# 输出数据集
dataiter = iter(trainloader)
images, labels = dataiter.next()
print(images.shape, labels.shape)

在这个示例中，我们首先定义了一个名为transform的Compose对象，其中包含了两个预处理函数：ToTensor和Normalize。然后，我们使用torchvision.datasets.CIFAR10函数加载CIFAR10数据集，并将transform对象传递给transform参数。最后，我们使用torch.utils.data.DataLoader函数加载数据集，并使用iter函数和next函数获取一个batch的数据。