作者: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
来识别数据集中的给定样本。一个跟踪样本和它们的标签的好方法就是采用下面的框架:
创建一个名为
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 torchclass 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 进行神经网络的计算。
译者注
往期回顾
一篇长文学懂入门推荐算法库:surprise
pytorch学习笔记(1):开始一个简单的分类器
pytorch学习笔记(2):在 MNIST 上实现一个 cnn
pytorch学习笔记(3):常用网络层介绍
pytorch学习笔记(4):tensorboard 可视化
pytorch学习笔记(5):vgg 实现以及一些 tricks
pytorch学习笔记(6):GPU 和如何保存加载模型
pytorch学习笔记(7):RNN 和 LSTM 实现分类和回归