超原版速度110倍，针对PyTorch的CPU到GPU张量迁移工具开源

选自Github

作者：SantoshGupta

机器之心编译

参与：杜伟、一鸣、泽南

神经网络的训练中往往需要进行很多环节的加速，这就是为什么我们逐渐使用GPU替代CPU、使用各种各样的算法来加速机器学习过程。但是，在很多情况下，GPU并不能完成CPU进行的很多操作。比如训练词嵌入时，计算是在CPU上进行的，然后需要将训练好的词嵌入转移到GPU上进行训练。

在这一过程中，张量迁移可能会变得非常缓慢，这成为了机器学习训练的一个瓶颈。

面对这样的问题，在很多优化方法中，就是将操作尽可能放在GPU上（如直接在GPU上进行数据预处理、词嵌入等的操作），并努力减少两者之间的数据交互，因为这些环节都很费时。机器之心就曾报道过这样的一个教程，将数据预处理放在了GPU上，减少了不少时间。

以上事例说明，如果能够做好CPU和GPU之间的迁移，则可以帮助开发者更好地优化机器学习模型，使CPU、GPU等硬件更好地完成自己的工作。

近日，有一位开发者就开源了一个名为SpeedTorch的工具。这一工具库可以实现高达110倍的CPU到GPU迁移加速。

项目地址：

项目背景

作者表示，最初想要创建SpeedTorch库是为了帮助训练大量的嵌入向量，而GPU在RAM保存这些嵌入时可能有困难。为了解决这个问题，他发现在CPU上托管一些此类嵌入有助于在GPU上保存它们。嵌入系统采用稀疏训练，只有一部分参数参与前馈/更新操作，剩余参数处于闲置状态。所以作者想到，为什么不在训练的过程中关闭这些闲置参数呢？这就需要快速的CPU→GPU数据迁移能力。

随着CPU→GPU迁移速度的加快，除了加速了CPU到GPU的张量转移外，开发者还可以实现很多新的功能。

将SpeedTorch库嵌入数据管道中，实现CPU和GPU之间快速的双向数据交互；

通过CPU存储将模型的训练参数增加近两倍（闲置参数保存在CPU中，需要更新时再移动到GPU里，因此可以扩大模型整体的参数量）；

在训练稀疏嵌入向量中采用Adadelta、Adamax、RMSprop、Rprop、ASGD、AdamW和Adam优化器。之前只有SpraseAdam、Adagrad和SGD适合稀疏嵌入训练。

那么，能够实现如此惊人的加速的库是怎么实现的呢？

SpeedTorch

背后的技术

SpeedTorch如此之快的技术是因为它是基于Cupy开发的。CuPy是一个借助CUDAGPU库在英伟达GPU上实现Numpy数组的库。基于Numpy数组的实现，GPU自身具有的多个CUDA核心可以促成更好的并行加速。

CuPy接口是Numpy的一个镜像，并且在大多情况下，它可以直接替换Numpy使用。只要用兼容的CuPy代码替换Numpy代码，用户就可以实现GPU加速。

CuPy支持Numpy的大多数数组运算，包括索引、广播、数组数学以及各种矩阵变换。

有了这样强大的底层支持，再加上一些优化方法，SpeedTorch就能达到110倍的速度了。

使用方法

SpeedTorch可以通过pip安装。你需要在导入SpeedTorch之前事先安装和导入Cupy。

安装步骤如下：

!pipinstallSpeedTorchimportcupyimportSpeedTorch

利用SpeedTorch加快CPU→GPU数据迁移速度

如下colabnotebook所示，如何利用DataGadget将数据载入SpeedTorch，以及如何将数据移入/移出Pytorchcuda变量。

代码示例：

借助于SpeedTorch将非稀疏优化器（本例中为Adamax）用于稀疏训练

SkipGram_ModelRegular=SkipGramModelRegular(numEmbeds=number_items,emb_dimension=128,sparseB=True)use_cuda=_available()ifuse_cuda:SkipGram_()optimizer=(SkipGram_())runningLoss=0runnngTime=0batch_size=512negSamp=64numPos=4skip_window=int(numPos/2)targets=(batch_size,numPos+negSamp,dtype=).cuda()foriinrange(500):batch,labels,negz=generate_batch(batch_size=batch_size,skip_window=skip_window,negRate=negSamp)batchTensor=_numpy(batch)LabelTensor=_numpy(labels)negTensor=_numpy(negz)pos_u=Variable((()))pos_v=Variable((()))neg_v=Variable((()))ifuse_cuda:pos_u=pos_()pos_v=pos_()neg_v=neg_()_grad()loss=SkipGram_(pos_u,pos_v,neg_v,targets)runningLoss=runningLoss+()()()

代码示例:

以上展示了如何以常规的方式训练word2vec，随后展示了如何使用SpeedTorch在同样的数据上进行训练——在通常不支持稀疏训练的优化器上。因为嵌入变量包含的所有嵌入在每一部上都有更新，你可以在初始化期间将sparse=False。

效果

这一部分记录了Cupy/PyTorch张量和PyTorch变量之间的数据迁移速度。其中，需要迁移128维的嵌入向量，共有131,072个32位浮点数。使用了如下的代码进行测试工作。所有测试都使用了特斯拉K80GPU。

测试代码链接：

下表是结果摘要。在同样情况下，将数据从PyTorchCUDA张量传递到CUDAPyTorch嵌入变量上是要比SpeedTorch更快的，但对于所有其他的传输类型，SpeedTorch更快。对于转移到CudaPytorch嵌入，或从CudaPytorch嵌入转移的两个步骤的总和上来说，SpeedTorch比常规GPU和CPUPinned张量的Pytorch速度同样快。

从表中可以看出，这是SpeedTorch确实比PyTorch自带的数据迁移方法要快很多。