Pytorch中的torch.utils.data.Dataset 类

1、使用方法

python 复制代码
from torch.utils.data import Dataset

2、torch.utils.data.Dataset 类的定义

python 复制代码
class Dataset(Generic[_T_co]):
    r"""An abstract class representing a :class:`Dataset`.

    All datasets that represent a map from keys to data samples should subclass
    it. All subclasses should overwrite :meth:`__getitem__`, supporting fetching a
    data sample for a given key. Subclasses could also optionally overwrite
    :meth:`__len__`, which is expected to return the size of the dataset by many
    :class:`~torch.utils.data.Sampler` implementations and the default options
    of :class:`~torch.utils.data.DataLoader`. Subclasses could also
    optionally implement :meth:`__getitems__`, for speedup batched samples
    loading. This method accepts list of indices of samples of batch and returns
    list of samples.

    .. note::
      :class:`~torch.utils.data.DataLoader` by default constructs an index
      sampler that yields integral indices.  To make it work with a map-style
      dataset with non-integral indices/keys, a custom sampler must be provided.
    """

    def __getitem__(self, index) -> _T_co:
        raise NotImplementedError("Subclasses of Dataset should implement __getitem__.")

    # def __getitems__(self, indices: List) -> List[_T_co]:
    # Not implemented to prevent false-positives in fetcher check in
    # torch.utils.data._utils.fetch._MapDatasetFetcher

    def __add__(self, other: "Dataset[_T_co]") -> "ConcatDataset[_T_co]":
        return ConcatDataset([self, other])

    # No `def __len__(self)` default?
    # See NOTE [ Lack of Default `__len__` in Python Abstract Base Classes ]
    # in pytorch/torch/utils/data/sampler.py

原文解释:

  1. 所有表示从键到数据样本映射的数据集都应该继承自这个类。

    这意味着,如果你有一个数据集,它通过某些键(可能是整数、字符串等)来访问数据样本,那么你应该从 Dataset 类继承来创建你的数据集类。

  2. 所有的子类都应该重写 __getitem__ 方法,支持通过给定的键获取数据样本。
    __getitem__ 是 Python 的特殊方法,用于通过 dataset[key] 这样的语法来获取数据。在你的子类中,你需要实现这个方法,确保它能够返回与给定键对应的数据样本。

  3. 子类也可以选择性地重写 __len__ 方法,该方法通常被许多 Sampler 实现和 DataLoader 的默认选项所使用,用于返回数据集的大小。
    __len__ 方法用于返回数据集中样本的总数。虽然它不是强制要求的,但如果你希望使用 PyTorch 的 SamplerDataLoader,通常需要实现这个方法。

  4. 子类还可以选择性地实现 __getitems__ 方法,以加速批量数据加载。这个方法接受一个包含批次样本索引的列表,并返回一个样本列表。
    __getitems__ 是一个可选的优化方法。如果你需要批量加载数据,实现这个方法可以提高效率。它接受一个索引列表,并返回对应的样本列表。

  5. DataLoader 默认构造一个生成整数索引的采样器(sampler)。要使它能够与具有非整数索引/键的 map-style 数据集一起工作,必须提供一个自定义的采样器。
    DataLoader 默认情况下假设你的数据集是可以通过整数索引访问的(即 dataset[0], dataset[1] 等)。如果你的数据集使用非整数键(比如字符串或其他类型),你需要提供一个自定义的采样器来生成这些键。

3、示例

示例 1:简单的整数索引数据集

假设我们有一个数据集,数据储存在一个列表中,我们可以通过整数索引来访问。

python 复制代码
from torch.utils.data import Dataset

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

    def __getitem__(self, index):
        return self.data[index]

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

# 示例数据
data = [1, 2, 3, 4, 5]

# 创建数据集
dataset = SimpleDataset(data)

# 使用
print(dataset[0])  # 输出:1
print(len(dataset))  # 输出:5

示例 2:字符串键的数据集

假设我们有一个数据集,数据以字典形式存储,键是字符串。

python 复制代码
from torch.utils.data import Dataset

class StringKeyDataset(Dataset):
    def __init__(self, data):
        self.data = data
        self.keys = list(data.keys())

    def __getitem__(self, key):
        return self.data[key]

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

# 示例数据
data = {"a": 1, "b": 2, "c": 3}

# 创建数据集
dataset = StringKeyDataset(data)

# 使用
print(dataset["a"])  # 输出:1
print(len(dataset))  # 输出:3

注意:如果需要与 DataLoader 一起使用,必须提供一个自定义的采样器,因为默认的采样器生成整数索引。

示例 3:实现 __getitems__ 方法

为了实现批量加载数据,我们可以实现 __getitems__ 方法。

python 复制代码
from torch.utils.data import Dataset

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

    def __getitem__(self, index):
        return self.data[index]

    def __getitems__(self, indices):
        return [self.data[i] for i in indices]

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

# 示例数据
data = [10, 20, 30, 40, 50]

# 创建数据集
dataset = BatchableDataset(data)

# 使用
print(dataset[0])  # 输出:10
print(dataset.__getitems__([1, 3]))  # 输出:[20, 40]

示例 4:图像数据集

假设我们有一个图像数据集,图像路径存储在列表中。

python 复制代码
from torch.utils.data import Dataset
from PIL import Image
import os

class ImageDataset(Dataset):
    def __init__(self, img_dir):
        self.img_dir = img_dir
        self.img_names = os.listdir(img_dir)

    def __getitem__(self, index):
        img_name = self.img_names[index]
        img_path = os.path.join(self.img_dir, img_name)
        image = Image.open(img_path)
        return image

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

# 创建数据集
img_dir = "path/to/images"
dataset = ImageDataset(img_dir)

# 使用
print(len(dataset))  # 输出图像数量
print(dataset[0])  # 输出第一张图像

示例 5:自定义采样器

如果你的数据集使用非整数键(如字符串),并且你想与 DataLoader 一起使用,可以定义一个自定义采样器。

python 复制代码
from torch.utils.data import Dataset, DataLoader, Sampler
import random

class StringKeyDataset(Dataset):
    def __init__(self, data):
        self.data = data
        self.keys = list(data.keys())

    def __getitem__(self, key):
        return self.data[key]

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

class StringSampler(Sampler):
    def __init__(self, keys):
        self.keys = keys

    #每次调用时(如新的epoch开始),先打乱键的顺序,再返回迭代器。
    #实现数据加载时的随机化顺序。
    def __iter__(self):
        random.shuffle(self.keys)
        return iter(self.keys)

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

# 示例数据
data = {"a": 1, "b": 2, "c": 3}

# 创建数据集和采样器
dataset = StringKeyDataset(data)
sampler = StringSampler(dataset.keys)

# 使用 DataLoader
dataloader = DataLoader(dataset, sampler=sampler, batch_size=2)

for batch in dataloader:
    print(batch)  # 输出批次数据
相关推荐
强哥之神14 分钟前
Meta AI 推出 Multi - SpatialMLLM:借助多模态大语言模型实现多帧空间理解
人工智能·深度学习·计算机视觉·语言模型·自然语言处理·llama
喜欢吃豆24 分钟前
从零构建MCP服务器:FastMCP实战指南
运维·服务器·人工智能·python·大模型·mcp
一个处女座的测试44 分钟前
Python语言+pytest框架+allure报告+log日志+yaml文件+mysql断言实现接口自动化框架
python·mysql·pytest
nananaij1 小时前
【Python基础入门 re模块实现正则表达式操作】
开发语言·python·正则表达式
蛋仔聊测试1 小时前
Playwright 网络流量监控与修改指南
python
nightunderblackcat2 小时前
进阶向:Python音频录制与分析系统详解,从原理到实践
开发语言·python·音视频
慕婉03072 小时前
Tensor自动微分
人工智能·pytorch·python
MUTA️2 小时前
pycharm中本地Docker添加解释器
ide·python·pycharm
神经星星2 小时前
专治AI审稿?论文暗藏好评提示词,谢赛宁呼吁关注AI时代科研伦理的演变
人工智能·深度学习·机器学习
想要成为计算机高手3 小时前
4. isaac sim4.2 教程-Core API-Hello robot
人工智能·python·机器人·英伟达·isaac sim·仿真环境