from torch.utils.data import Dataset详解-CSDN博客

本文链接：https://blog.csdn.net/m0_49133355/article/details/140553781

torch.utils.data.Dataset 是 PyTorch 数据加载库中的一个重要类，用于定义自定义数据集。通过继承 Dataset 类，可以创建自己的数据集类，并实现数据的加载和处理逻辑。下面是对 torch.utils.data.Dataset 的详细介绍：

基本概念
Dataset 类是一个抽象类，需要用户继承并实现以下两个方法：

len(self)：返回数据集的大小。
getitem(self, idx)：支持索引操作，返回指定索引的样本。
2. 使用步骤
2.1 导入相关模块

import torch
from torch.utils.data import Dataset

2.2 定义自定义数据集类
继承 Dataset 类，并实现 len 和 getitem 方法。

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

2.3 使用自定义数据集

# 示例数据
data = torch.randn(100, 3, 224, 224)  # 100个样本，每个样本为3通道224x224图像
labels = torch.randint(0, 10, (100,))  # 100个样本的标签，标签范围为0到9

# 创建自定义数据集实例
dataset = CustomDataset(data, labels)

# 访问数据集中的元素
print(len(dataset))  # 输出: 100
print(dataset[0])    # 输出: (tensor, label) 形式的数据和标签

与 DataLoader 配合使用
通常，Dataset 类与 torch.utils.data.DataLoader 配合使用，以便批量加载数据、打乱数据、并行加速等。

from torch.utils.data import DataLoader

# 创建DataLoader实例
dataloader = DataLoader(dataset, batch_size=4, shuffle=True)

# 迭代加载数据
for batch_data, batch_labels in dataloader:
    print(batch_data.size(), batch_labels.size())

高级用法
4.1 数据预处理
在 getitem 方法中添加数据预处理步骤。

from torchvision import transforms

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

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

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

# 定义数据预处理
transform = transforms.Compose([
    transforms.ToPILImage(),
    transforms.Resize((128, 128)),
    transforms.ToTensor(),
])

dataset = CustomDataset(data, labels, transform=transform)

4.2 从文件加载数据
如果数据存储在文件中，可以在 init 方法中读取文件路径，并在 getitem 方法中加载文件数据。

import pandas as pd

class CustomDataset(Dataset):
    def __init__(self, csv_file):
        self.data_frame = pd.read_csv(csv_file)

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

    def __getitem__(self, idx):
        row = self.data_frame.iloc[idx]
        sample = torch.tensor(row['data'], dtype=torch.float32)
        label = torch.tensor(row['label'], dtype=torch.long)
        return sample, label

dataset = CustomDataset('data.csv')