作者：Afshine Amidi ， Shervine Amidi

编译：silver

分享一篇斯坦福的两位同胞大佬的文章，这两位大佬的很多文章被机器之心等大号多次转载，他们的 gayhub 也被多次介绍。这次偶然看到一篇他们的文章，刚好最近在写 pytorch 的笔记，就分享过来，大家一起动手试试吧~

以下是全文：

动机

是否曾经有那么一刻，你不得不消耗大量的内存资源来读取数据，然后你希望能有一个魔术戏法来无缝的解决这一切？随着我们可以处理的数据量不断增加，大规模的数据集也逐渐成为我们生活的一部分。

我们必须承认，在某些情况下，即使最先进的计算机配置，也没有足够的内存空间按照我们过去的方式来处理数据。这也就是为什么我们需要去寻找一个有效的方法来完成这个任务。在这篇文章中，我们将会展示怎样实时的在多核上生成你的数据，并且把它立刻喂给你的深度学习模型。

本教程会向你展示如何在 pytorch 框架上实现这个功能。在训练数据的过程中，有效的数据生成方式对充分利用 GPU 的潜能是至关重要的。

教程

以前的方法

在开始正文之前，回忆一下你以前的 pytorch 代码是不是像下面这样：

# Load entire datasetX, y = torch.load('some_training_set_with_labels.pt')
# Train modelfor epoch in range(max_epochs):    for i in range(n_batches):        # Local batches and labels        local_X, local_y = X[i*n_batches:(i+1)*n_batches,], y[i*n_batches:(i+1)*n_batches,]
        # Your model        [...]

或者是这样：

# Unoptimized generatortraining_generator = SomeSingleCoreGenerator('some_training_set_with_labels.pt')
# Train modelfor epoch in range(max_epochs):    for local_X, local_y in training_generator:        # Your model        [...]

本教程是关于优化整个数据生成的过程，因此这些将不会再成为你训练程序的瓶颈。

为了便于理解，我们将这种情景下构建并行式的数据生成的研究逐步分割。BTW，下面介绍的代码都是可以用到你的项目中的很好的框架代码，你可以直接复制粘贴下面的代码块，然后对应的把空缺补上即可。

标记

在开始之前，让我们先整理一些在处理大规模数据集时特别有用的组织技巧。

用 python 字符串 ID 来识别数据集中的给定样本。一个跟踪样本和它们的标签的好方法就是采用下面的框架：

1. 创建一个名为 partition 的字典包含如下信息：

• 在 partition['train'] 中是一个包含了训练集 ID 的 list

• 在 partition['validation']中是一个包含了验证集 ID 的 list

• 创建一个名为 labels 的字典，其中包含了数据集中的每个 ID，由 labels[ID] 来实现数据和标签的关联

举个栗子，让我们假设有一个训练集包含了 id-1，id-2 和 id-3，对应的标签分别为 0，1 和 2，而验证集中包含了 id-4 和它的标签 1。在这个栗子中，用 python 定义的变量 partition 和 labels 就是这样

>>> partition{'train': ['id-1', 'id-2', 'id-3'], 'validation': ['id-4']}

和

>>> labels{'id-1': 0, 'id-2': 1, 'id-3': 2, 'id-4': 1}

另外，为了模块化，我们会分开实现 PyTorch 代码和自定义类，因此你的文件看起来会是这样：

folder/├── my_classes.py├── pytorch_script.py└── data/

其中 data/ 是包含了你的数据集的文件夹。

最后，值得高兴的一点是，本教程中的代码旨在使其通用化和最小化，因此你可以很容易的将其应用到你自己的数据集上。

Dataset 类

现在，让我们进入到如何构建 Python 类 Dataset 的细节中，它将描述你想要生成的数据集的关键特征。

首先，让我们写出这个类的初始化函数。我们让其继承 torch.utils.data.Dataset 类，这样我们就可以在后面利用一些诸如 多线程 之类很棒的功能。

def __init__(self, list_IDs, labels):    'Initialization'    self.labels = labels    self.list_IDs = list_IDs

这里，我们存储了一些重要的信息，例如 labels 的列表和 IDs 的列表这些我们希望在每一步来生成的内容。

每次调用都获取一个样本的索引，其上界由 __len__ 方法确定。

def __len__(self):    'Denotes the total number of samples'    return len(self.list_IDs)

现在，当通过一个给定索引来调用对应的样本时，生成器通过执行 __getitem__ 方法来生成它。

def __getitem__(self, index):    'Generates one sample of data'    # Select sample    ID = self.list_IDs[index]
    # Load data and get label    X = torch.load('data/' + ID + '.pt')    y = self.labels[ID]
    return X, y

在数据生成期间，这个方法从对应的文件 ID.pt 中读取给定例子的 Torch 张量。因为我们的代码是从易于多核处理的角度设计的，所以你可以做一些更复杂的操作进行代替（例如从源文件中进行计算），而无需担心数据生成会成为你训练过程的瓶颈。

这里对应我们本节描述内容的每一步的完整代码如下。

import torch
class Dataset(torch.utils.data.Dataset):  'Characterizes a dataset for PyTorch'  def __init__(self, list_IDs, labels):        'Initialization'        self.labels = labels        self.list_IDs = list_IDs
  def __len__(self):        'Denotes the total number of samples'        return len(self.list_IDs)
  def __getitem__(self, index):        'Generates one sample of data'        # Select sample        ID = self.list_IDs[index]
        # Load data and get label        X = torch.load('data/' + ID + '.pt')        y = self.labels[ID]
        return X, y

PyTorch 脚本

现在，我们需要对应地调整一下我们的 PyTorch 脚本，以便于它可以使用我们刚刚创建的生成器。为了实现这一点，我们使用了 PyTorch 的 DataLoader 类，这样除了我们自己创建的 Dataset 类，还可以囊括下面这些重要的参数：

• batch_size，这个参数表示了每批生成的数据包含样本的数量。

• shuffle，如果设置为 True，我们将会每次得到一个乱序的生成结果（反之则会线性顺序生成）。将喂给分类器的数据打乱是很好的做法，这样在不同的 epoch 中每个 batch 的数据不会看起来都一样。这种做法最终会使得我们的模型更鲁棒。

• num_workers，这个参数描述了并行生成批数据的进程数量。足够多的进程数可以确保有效管理 CPU 的计算性能，也就是说计算瓶颈会是在 GPU 上神经网络的前向传播和反向传播操作（而不会是数据生成部分）。

一个可以直接写在你的脚本中的代码模板建议如下所示。

import torchfrom my_classes import Dataset

# CUDA for PyTorchuse_cuda = torch.cuda.is_available()device = torch.device("cuda:0" if use_cuda else "cpu")torch.backends.cudnn.benchmark = True
# Parametersparams = {'batch_size': 64,          'shuffle': True,          'num_workers': 6}max_epochs = 100
# Datasetspartition = # IDslabels = # Labels
# Generatorstraining_set = Dataset(partition['train'], labels)training_generator = torch.utils.data.DataLoader(training_set, **params)
validation_set = Dataset(partition['validation'], labels)validation_generator = torch.utils.data.DataLoader(validation_set, **params)
# Loop over epochsfor epoch in range(max_epochs):    # Training    for local_batch, local_labels in training_generator:        # Transfer to GPU        local_batch, local_labels = local_batch.to(device), local_labels.to(device)
        # Model computations        [...]
    # Validation    with torch.set_grad_enabled(False):        for local_batch, local_labels in validation_generator:            # Transfer to GPU            local_batch, local_labels = local_batch.to(device), local_labels.to(device)
            # Model computations            [...]



总结


就是这样！现在你可以通过下面的命令行来运行你的 PyTorch 脚本了。

python3 pytorch_script.py

然后你将看到在训练阶段，数据是 CPU 并行生成的，然后可以被喂给 GPU 进行神经网络的计算。

译者注

原文地址：//stanford.edu/~shervine/blog/pytorch-how-to-generate-data-parallel

完整代码及示例数据集：//github.com/shervinea/pytorch-data-generator/