使用 Pytorch 数据加载器加载特定样本的简单方法

Question

我目前正在训练一个 3D CNN，用于具有相对稀疏标签的二进制分类（标签数据中约 1% 的体素对应于目标类）。

为了在训练期间执行基本的健全性检查（例如，网络是否完全学习？）向网络展示一小部分精心挑选的训练示例子集，其中目标类标签的比例高于平均水平，这会很方便.

按照 Pytorch 文档的建议，我实现了自己的 dataset 类（继承自 torch.utils.data.Dataset），它通过 __get_item__ 方法向 torch.utils.data.DataLoader 提供训练示例。

在我找到的pytorch tutorials 中，DataLoader 用作迭代器来生成训练循环，如下所示：

for i, data in enumerate(self.dataloader):

    # Get training data
    inputs, labels = data

    # Train the network
    # [...]

我现在想知道的是是否存在一种简单的方法来加载单个或几个特定的​​训练示例（使用 Dataset 的 __get_item__ 方法理解的线性索引）。但是，DataLoader 没有 __get_item__ 方法并反复调用 __next__ 直到我达到所需的索引似乎并不优雅。

显然，解决此问题的一种可能方法是定义一个自定义 sampler 或 batch_sampler，继承自抽象 torch.utils.data.Sampler。但这似乎在顶部检索一些特定的样本。

我想我在这里忽略了一些非常简单和明显的东西。任何建议表示赞赏！

Answer 1

以防万一有类似问题的人在某个时候遇到这个问题：

我最终使用的快速而简单的解决方法是绕过训练循环中的dataloader，直接访问它的关联dataset 属性。假设我们想快速检查我们的网络是否完全学习，方法是重复向它展示一个精心挑选的线性索引sample_idx（由数据集类定义）的训练示例。

然后可以做这样的事情：

for i, _ in enumerate(self.dataloader):

    # Get training data
    # inputs, labels = data

    inputs, labels = self.dataloader.dataset[sample_idx]
    inputs = inputs.unsqueeze(0)
    labels = labels.unsqueeze(0)

    # Train the network
    # [...]

编辑：

一个简短的评论，因为有些人似乎发现这个解决方法很有帮助：当使用这个黑客时，我发现用num_workers = 0 实例化DataLoader 是至关重要的。否则，可能会出现内存分割错误，在这种情况下，您最终可能会得到看起来非常奇怪的训练数据。

Answer 2

如果你已经定义了

train_set = torchvision.datasets.CIFAR10(root='~/datasets/', train=True,
                                                download=True, transform=(transform['train']))

然后你可以做类似的事情

train_set.data[index] 其中index 是您想要的具体示例的index。

现在，您可以使用包含这些特定示例的新数据集重新定义您的 Dataset 类，并且您拥有它。