SpeedTorch

在特定情况下，实现更快的固定 CPU 张量与 GPU PyTorch 变量之间的传输，以及 GPU 张量与 GPU PyTorch 变量之间的传输。

更新 9-29-19

由于在某些系统上，使用固定的 PyTorch CPU 张量比使用 Cupy 张量更快（详情请参阅“工作原理”部分），我创建了通用的 PyTorch 张量类 PytorchModelFactory 和 PytorchOptimizerFactory，它们可以指定将张量设置为 cuda 或 cpu，并且如果使用 cpu，还可以选择是否将其内存固定。原有的 GPUPytorchModelFactory 和 GPUPytorchOptimizerFactory 类仍然保留在库中，因此现有使用 SpeedTorch 的代码不会受到影响。文档也已更新，加入了这些新类。

它是什么？

该库的核心是固定在 CPU 上的 Cupy 张量，相比普通的 PyTorch 固定 CPU 张量，它可以使 CPU 到 GPU 的传输速度提升 3.1 倍，而 GPU 到 CPU 的传输速度则可提升至 410 倍。具体速度取决于数据量和系统中的 CPU 核心数（更多细节请参阅“工作原理”部分）。

该库还包含用于嵌入训练的函数；当嵌入处于空闲状态时，可以将其存储在 CPU 内存中，从而节省 GPU 内存。

灵感来源

我最初创建这个库是为了帮助训练大量嵌入，而 GPU 可能难以容纳这些数据。为此，我发现将部分嵌入存储在 CPU 上可以帮助实现这一目标。嵌入系统采用稀疏训练方式——只有总参数的一小部分参与前向传播或更新步骤，其余部分则处于空闲状态。因此，我想：为什么不在训练过程中将空闲参数从 GPU 上移开呢？这就需要快速的 CPU 到 GPU 数据传输。

如需了解完整背景，请访问 Devpost 页面：

https://devpost.com/software/speedtorch-6w5unb

快速的 CPU 到 GPU 传输能为我带来什么？（远超你的想象）

借助快速的 CPU 到 GPU 传输，我们可以开发出许多有趣的方法来实现以前被认为不可能的功能。

🏎️ 将 SpeedTorch 集成到你的数据管道中，实现 CPU 与 GPU 之间快速的数据传输。

🏎️ 通过 CPU 存储扩充训练参数。只要你的 CPU 内存足够，就可以存储任意数量的嵌入，而无需担心 GPU 内存不足的问题。

🏎️ 使用 Adadelta、Adamax、RMSprop、Rprop、ASGD、AdamW 和 Adam 优化器进行稀疏嵌入训练。过去，只有 SparseAdam、Adagrad 和 SGD 适合处理稀疏梯度，因为只有这些优化器直接支持稀疏梯度。

基准测试

速度

（编辑于 9-20-19，一位 PyTorch 开发者指出了原始基准测试代码中的一些小错误，数值和代码均已更新）

这里有一个笔记本，比较了通过 SpeedTorch 和 PyTorch 张量进行传输的效果，包括固定 CPU 和 CUDA 张量两种情况。所有测试均在配备 Tesla K80 GPU 和 2 核 CPU 的 Colab 实例上完成。

更新 10-17-19：Google Colab 现在默认配备 4 核 CPU，因此该笔记本的结果会与下方报告的不同，因为随着 CPU 核心数的增加，PyTorch 的索引内核效率也会提高。CPU 核心数更多的系统在使用 SpeedTorch 时优势会相对减小。

https://colab.research.google.com/drive/1PXhbmBZqtiq_NlfgUIaNpf_MfpiQSKKs

此笔记本测量了 131,072 个维度为 128 的 float32 嵌入，在 Cupy/PyTorch 张量和 PyTorch 变量之间来回传输的时间，重复次数为 n=100。Google Colab 的 CPU 是 4 核的，这会影响传输速度。CPU 核心数较多的系统使用 SpeedTorch 的优势会较小。

下表总结了测试结果。从 PyTorch CUDA 张量传输数据到 CUDA PyTorch 嵌入变量的速度比 SpeedTorch 对应操作更快，但在其他所有传输类型中，SpeedTorch 都更快。而对于往返 CUDA PyTorch 嵌入变量的总时间，无论是普通 GPU 张量还是固定 CPU 张量，SpeedTorch 都比 PyTorch 对应操作更快。

我注意到不同 Colab 实例的运行结果可能会有所不同，因此在查看这些结果时请注意这一点。亲自运行该 Colab 笔记本可能会得到不同的数值，不过结果的数量级通常是一致的。

以下表格中的传输时间单位为秒。本次基准测试是在 CPU 为 2 核的 Colab 实例上进行的。Colab 还有付费的 Pro 版本，配备 4 核 CPU，因此以下基准测试结果并不适用于那些实例。

张量类型	传输至 CUDA PyTorch 变量	对比
SpeedTorch(cuda)	0.0087	比 PyTorch 对应操作慢 6.2 倍
SpeedTorch(PinnedCPU)	0.0154	比 PyTorch 对应操作快 3.1 倍
PyTorch(cuda)	0.0014	比 SpeedTorch 对应操作快 6.2 倍
PyTorch(PinnedCPU)	0.0478	比 SpeedTorch 对应操作慢 3.1 倍

张量类型	从 CUDA PyTorch 变量传输	对比
SpeedTorch(cuda)	0.0035	比 PyTorch 对应操作快 9.7 倍
SpeedTorch(PinnedCPU)	0.0065	比 PyTorch 对应操作快 410 倍
PyTorch(cuda)	0.0341	比 SpeedTorch 对应操作慢 9.7 倍
PyTorch(PinnedCPU)	2.6641	比 SpeedTorch 对应操作慢 410 倍

张量类型	往返 CUDA PyTorch 变量的总时间	对比
SpeedTorch(cuda)	0.0122	比 PyTorch 对应操作快 2.9 倍
SpeedTorch(PinnedCPU)	0.0219	比 PyTorch 对应操作快 124 倍
PyTorch(cuda)	0.0355	比 SpeedTorch 对应操作慢 2.9 倍
PyTorch(PinnedCPU)	2.7119	比 SpeedTorch 对应操作慢 124 倍

类似的基准测试也针对 PyTorch CUDA 优化器的传输进行了计算。结果基本相同，以下是用于优化器基准测试的笔记本：

https://colab.research.google.com/drive/1Y2nehd8Xj-ixfjkj2QWuA_UjQjBBHhJ5

内存

尽管 SpeedTorch 的张量通常比 PyTorch 的更快，但缺点是 SpeedTorch 的张量会占用更多内存。不过，由于数据传输速度更快，你可以利用 SpeedTorch 在 GPU 和 CPU 上同时保存参数，从而增加架构中训练的嵌入数量。

下表是在 Google Colab 中进行基准测试的汇总。根据我的经验，Colab 报告的内存值似乎存在一些波动，大约在 ±0.30 GB 左右，因此在查看这些数值时请注意这一点。这些数值是针对一个 10,000,000×128 的 float32 张量而言的。

张量类型	CPU (GB)	GPU (GB)
Cupy PinnedCPU	9.93	0.06
Pytorch PinnedCPU	6.59	0.32
Cupy Cuda	0.39	9.61
Pytorch Cuda	1.82	5.09

虽然 PyTorch GPU 张量与 PyTorch CUDA 变量之间的传输时间不如 Cupy 的对应操作快，但速度仍然可以接受。因此，如果内存仍然是一个问题，一种折中的方法是使用 SpeedTorch 的 Cupy CPU 固定张量来在 CPU 上存储参数，而使用 SpeedTorch 的 PyTorch GPU 张量来在 GPU 上存储参数。

这是我用来测量每种变量类型占用多少内存的笔记本： https://colab.research.google.com/drive/1ZKY7PyuPAIDrnx2HdtbujWo8JuY0XkuE 如果在 Colab 中使用此笔记本，每次创建完一个张量后都需要重启环境，以便为下一个张量获取准确的测量结果。

哪些系统能获得速度优势？

对于 CPU↔GPU 数据传输，这取决于传输的数据量以及你拥有的核心数。一般来说，对于 1–2 个 CPU 核心，SpeedTorch 会快得多。但随着 CPU 核心数的增加，PyTorch 的 CPU↔GPU 索引操作效率会更高。有关这方面的详细信息，请参阅下一节“工作原理”。如果你想简单地判断自己的系统是否能从中受益，可以在本地运行基准测试代码，并将数据量调整为你实际应用中会处理的数量。

对于 GPU↔GPU 数据传输，如果使用原生 PyTorch 的常规索引方式，所有系统都会获得速度提升，因为 SpeedTorch 绕过了 PyTorch 索引操作中的一个 bug。不过，这个 bug 可以通过使用 nightly 版本或改用不同的索引方式来避免，详情请参阅“工作原理”部分。

如何工作？

更新：2019 年 9 月 20 日——起初我并不清楚为什么 SpeedTorch 比 PyTorch 张量更快；我只是偶然发现了这一速度优势。后来，PyTorch 论坛上的一位开发者指出了其中的原因。

至于更好的 CPU↔GPU 传输，这是因为 SpeedTorch 通过将 CPU 张量伪装成 GPU 张量来避免一次 CPU 索引操作。如果 CPU 核心较少（例如 Google Colab 中只有 2 个核心），CPU 索引操作可能会很慢；但如果核心较多，则可能更快。具体效果取决于你要传输的数据量以及你的核心数量。

至于更好的 GPU↔GPU 传输，这是因为 SpeedTorch 避免了索引操作中的一个 bug。这个 bug 也可以通过使用 nightly 构建版本，或者在 PyTorch 1.1/1.2 中使用 index_select 或 index_copy_ 而不是 a[idx] 语法来避免。

有关详细信息，请参阅以下 PyTorch 帖子：

https://discuss.pytorch.org/t/introducing-speedtorch-4x-speed-cpu-gpu-transfer-110x-gpu-cpu-transfer/56147/2

其中一位 PyTorch 工程师详细分析了 Cupy 索引内核在某些情况下如何带来速度提升。实际上，加快速度的并不是传输本身，而是所使用的索引内核。

关于 Cupy 的内存管理机制，我建议参考我在 Stack Overflow 上提出的两个问题。聪明的用户 Robert Crovella 不仅给出了详细的解释，还通过开发自己的 Cupy 内存分配器，找到了为 Cupy 数组分配固定内存的方法。这基本上就是 SpeedTorch 背后的核心技术。

https://stackoverflow.com/questions/57750125/cupy-outofmemoryerror-when-trying-to-cupy-load-larger-dimension-npy-files-in-me https://stackoverflow.com/questions/57752516/how-to-use-cuda-pinned-zero-copy-memory-for-a-memory-mapped-file

使用指南

开始使用

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

!pip install SpeedTorch
import cupy
import SpeedTorch

使用 SpeedTorch 提高 CPU 到 GPU 的数据传输速度

这个 Colab 笔记本展示了如何使用 SpeedTorch 的 Data Gadget 将数据加载到 SpeedTorch 中，以及如何将这些数据与 PyTorch 的 CUDA 变量进行双向传输。

https://colab.research.google.com/drive/185Z5Gi62AZxh-EeMfrTtjqxEifHOBXxF

请参阅速度基准测试笔记本，以了解使用 SpeedTorch 所带来的速度优势。

使用 SpeedTorch 对稀疏数据使用非稀疏优化器（如 Adamax）

对于初次尝试使用 SpeedTorch 的人，我建议按照这个示例操作，因为 Word2Vec 是机器学习中较为常见的算法之一。

https://colab.research.google.com/drive/1ApJR3onbgQWM3FBcBKMvwaGXIDXlDXOt

该笔记本首先展示了如何以常规方式训练 Word2Vec，然后演示了如何使用 SpeedTorch 在相同数据上进行训练，采用一种通常不支持稀疏训练的优化器。之所以可行，是因为嵌入变量中的所有嵌入在每一步都会被更新，因此在初始化时可以将 sparse 参数设置为 False。

通过 CPU 存储扩充训练参数

简而言之：

常规训练：PyTorch 的嵌入变量包含所有嵌入；PyTorch 优化器则包含每个嵌入对应的所有参数权重。

SpeedTorch 训练：PyTorch 的嵌入变量仅包含一批嵌入；PyTorch 优化器也只包含该批次对应的参数权重。其余部分由 SparseTorch 张量存储，并在每一步与 PyTorch 变量交换嵌入或权重。

在诸如 word2vec、GloVe 或神经协同过滤等稀疏训练算法中，每一步仅训练总参数（嵌入）的一小部分。如果您的 GPU 无法以期望的嵌入维度容纳所有嵌入，一种选择是将部分参数托管在固定在内存中的 CPU Cupy 数组上，并在需要时将其传输到模型张量中。如果主要在 PyTorch 中执行此操作会非常缓慢，尤其是当在挂载了 CUDA 的 PyTorch 变量和固定在内存中的 CPU PyTorch 张量之间传输参数时，可能需要 2.5 到 3 秒（在 Google Colab 上）。幸运的是，使用 SpeedTorch 这一步仅需 0.02 到 0.03 秒！

应用场景：

--2,829,853 本书的嵌入--

SpeedTorch 被用于为一本稀有书籍推荐系统训练 2,829,853 本书的嵌入。

https://github.com/Santosh-Gupta/Lit2Vec2

https://devpost.com/software/lit2vec2

每本书的嵌入维度为 400，但理论上可以使用 496 维的嵌入。之所以采用 400 维，是因为我存储训练后嵌入的 Google Drive 空间有限 :( 。不过，现在 GPU 内存的限制不再是问题了 :) 以下是一个演示训练笔记本的直接链接，它使用 SpeedTorch 以 496 维嵌入进行训练。

注意：您需要使用具有 25 GB RAM 的 Colab 笔记本版本，而不是通常的 12 GB。要获得这种类型的实例，您需要先让当前实例因内存不足而崩溃，随后左下角会出现提示，询问您是否希望升级。您可以通过编写一个不断将 NumPy 浮点矩阵大小翻倍的循环来实现这一点。

https://colab.research.google.com/drive/1AqhT-HetihXMET1wJQROrC3Q9tFJqJ19

以下是使用相同模型和数据，但未使用 SpeedTorch 的直接链接：

https://colab.research.google.com/drive/1idV1jBOUZVPCfdsy40wIrRPHeDOanti_

采用传统训练方法时，Colab 能够处理的最大嵌入维度为 255–260，超过这个值就会出现 CUDA 内存不足错误：

RuntimeError: CUDA out of memory. Tried to allocate 2.74 GiB (GPU 0; 11.17 GiB total capacity; 8.22 GiB already allocated; 2.62 GiB free; 5.05 MiB cached)

--14,886,544 篇研究论文的嵌入--

https://github.com/Santosh-Gupta/Research2Vec2

借助 SpeedTorch，我可以在 GPU 上保留上下文嵌入的同时，将目标嵌入存储在 CPU 上，从而以 188 维的嵌入维度训练 14,886,544 篇研究论文的嵌入（由于使用了 SGD 优化器，因此没有优化器权重）。

以下是该笔记本的直接链接：

https://colab.research.google.com/drive/1saKzsaHoy6O_U1DF_z15_Qkr5YLNI_GR

注意：您需要使用具有 25 GB RAM 的 Colab 笔记本版本，而不是通常的 12 GB。要获得这种类型的实例，您需要先让当前实例因内存不足而崩溃，随后左下角会出现提示，询问您是否希望升级。您可以通过编写一个不断将 NumPy 浮点矩阵大小翻倍的循环来实现这一点。

如果没有 SpeedTorch，在 Google Colab 的 Tesla K80 GPU 上，嵌入维度只能达到 94–96，否则会出现 RuntimeError: CUDA out of memory 错误。以下是未使用 SpeedTorch 的训练版本：

https://colab.research.google.com/drive/1jh7RUgeajhdWdGNfWG3Twm1ZjyTQU0KR

最佳实践

1. 使用 Cupy GPU 张量时，请务必先初始化它们，再初始化固定在内存中的 CPU 张量。这是因为初始化 Cupy GPU 张量似乎会占用大量 CPU 内存。因此，如果您 CPU 内存有限，且已将固定在内存中的 CPU 张量加载到内存中，则初始化 Cupy GPU 张量可能会导致崩溃。

2. 如果您的所有参数都能装入 GPU 内存，请使用纯 PyTorch，因为这是最快的训练方式。但如果无法将所有参数装入内存，则应将参数（请注意，优化器也有权重）在 SpeedTorch 的 Cupy CUDA 张量和 Cupy 固定在内存中的 CPU 张量之间分配；这是第二快的方式。然而，即便如此仍无法将所有参数装入内存，则可将参数在 SpeedTorch 的 Cupy 固定在内存中的 CPU 张量和 SpeedTorch 的 PyTorch CUDA 张量之间分配；这种方式虽然比前两种慢，但使用的 GPU 内存更少。关于第三种方式，这里有两个示例笔记本：https://colab.research.google.com/drive/1AqhT-HetihXMET1wJQROrC3Q9tFJqJ19 和 https://colab.research.google.com/drive/1saKzsaHoy6O_U1DF_z15_Qkr5YLNI_GR。

3. 训练完成后，保存任何 CUDA 变量都会增加内存使用量，甚至可能导致崩溃，尤其是在使用 Cupy 时。如果您已经接近内存上限，建议使用 getNumpyVersion 方法获取张量的 NumPy 版本，然后使用 numpy.save 或 hdpy/pytables 保存您的 NumPy 数组。NumPy 保存更为轻量。

需要帮助？

您可以通过在 Gitter 上打开一个问题，或直接与我聊天：https://gitter.im/SpeedTorch。

未来工作

我正在考虑加入更多关于快速 CPU 到 GPU 数据传输的功能。如果您有任何想法，请在 GitHub 上提交一个问题。

实验性功能

除了 Cupy GPU/固定在内存中的 CPU 张量以及 PyTorch CUDA 张量之外，SpeedTorch 还提供 PyTorch 固定在内存中的 CPU 张量和 Cupy memmap GPU/固定在内存中的 CPU 张量。目前我尚未找到这些类型张量的实际用途，但它们已经完全实现并可供使用。

https://github.com/Santosh-Gupta/SpeedTorch/tree/master/SpeedTorch

我感兴趣的一个方向是探索是否可以通过使用 Cupy Memmaps 来减少内存占用。到目前为止，它们的内存占用与实时版本并无不同。

文档

类 ModelFactory

ModelFactory(model_variable,  total_classes,  embed_dimension, datatype = 'float32', CPUPinn = False)

使用 Cupy 为模型变量创建切换器。该切换器可以在完整的嵌入集合和模型批次集合之间切换变量。每个变量都需要一个独立的切换器。

示例：

uEmbed_switcher = SpeedTorch.ModelFactory( skip_gram_modelSparse.u_embeddings, total_classes=50000, embed_dimension=128)

参数：

model_variable：您希望为其创建切换器的模型中的特定变量。

total_classes：需要训练的嵌入总数。

embed_dimension：嵌入的维度。

datatype（可选）：变量的数据类型。默认为 'float32'。

CPUPinn（可选）：将完整的嵌入集合固定在 CPU 上。这样可以节省 GPU 内存，但数据传输会变慢。默认为 False。

方法：

zerosInit()：将变量切换器的完整集合初始化为零。

uniformDistributionInit(low, high)：将变量切换器的完整集合初始化为从 low 到 high 的均匀分布。

normalDistributionInit(mean, stdDev)：将变量切换器的完整集合初始化为均值为 mean、标准差为 stdDev 的正态分布。

variableTransformer( batchSize, posPerBatch, negPerBatch = None )：设置用于模型前向传播步骤的占位输入。batchSize 是批次大小，posPerBatch 是每批次的正样本数量。如果还需要用于负样本的第二个占位输入，则可以将 negPerBatch（可选）设置为负样本数量，此时将返回两个占位输入，而不是一个。

beforeForwardPass(retrievedPosIndexes , retrievedNegIndexes = None)：将嵌入从完整嵌入集合切换到模型嵌入集合。retrievedPosIndexes 是要检索的正样本索引。如果也需要检索负样本，则可以传递 retrievedNegIndexes（可选）作为负样本索引。

afterOptimizerStep( retrievedPosIndexes , retrievedNegIndexes = None)：将更新后的嵌入从模型切换回完整嵌入集合。retrievedPosIndexes 是之前检索的正样本索引。如果也检索了负样本，则可以传递 retrievedNegIndexes（可选）作为负样本索引。

saveCupy(saveFileName)：将张量保存为 .npy 文件。

loadCupy(loadFileName)：从 .npy 文件加载张量。

getNumpyVersion：获取张量的 NumPy 版本。

类 OptimizerFactory

OptimizerFactory( given_optimizer,  total_classes,  embed_dimension, model, variable_name, dtype='float32' , CPUPinn = False)

使用 Cupy 为优化器变量创建切换器。该切换器可以在完整的嵌入集合和优化器批次集合之间切换变量。每个变量都需要一个独立的切换器。

示例：

uAdagrad_switcher = SpeedTorch.OptimizerFactory(given_optimizer,  total_classes,  embed_dimension, model, variable_name, dtype='float32', CPUPinn = False)

参数：

given_optimizer：已用您的模型权重初始化的优化器。如果用于嵌入训练，请确保将 sparse 参数设置为 False。目前支持的优化器包括 SparseAdam、Adadelta、Adamax、Adam、AdamW、ASGD 和 RMSprop。Rprop 也被支持，但需要先完成第一次前向传播和 loss.backward() 步骤，才能初始化 OptimizerFactory 实例。这是因为 Rprop 优化器需要其参数的梯度来进行初始化。

total_classes：需要训练的嵌入总数。

embed_dimension：嵌入的维度。

model：您的模型实例。

variable_name：模型中定义的变量的确切名称。

dtype（可选）：您的变量的数据类型。默认为 'float32'。

CPUPinn（可选）：将完整的优化器变量权重集合固定在 CPU 上。这样可以节省 GPU 内存，但数据传输会变慢。默认为 False。

方法：

optInit：初始化优化器变量切换器。

beforeForwardPass(retrievedPosIndexes , retrievedNegIndexes = None)：将优化器变量权重从完整权重集合切换到优化器权重张量。retrievedPosIndexes 是要检索的正样本索引。如果也需要检索负样本，则可以传递 retrievedNegIndexes（可选）作为负样本索引。

afterOptimizerStep( retrievedPosIndexes , retrievedNegIndexes = None)：将优化器变量权重从您的优化器切换回完整权重集合。retrievedPosIndexes 是之前检索的正样本索引。如果也检索了负样本，则可以传递 retrievedNegIndexes（可选）作为负样本索引。

类 DataGadget

创建一个张量，其主要功能是将其内容传输到 PyTorch 的 CUDA 变量中。

DataGadget( fileName, CPUPinn=False)

参数：

fileName：要打开的 .npy 数据文件的位置。

CPUPinn（可选）：将您的数据固定在 CPU 上。默认为 False。

方法：

getData(indexes)：以 PyTorch Cuda 变量可以接受的格式检索数据。indexes 是要从中检索数据的张量索引。

insertData(dataObject, indexes)：从 PyTorch Cuda 变量中插入数据。dataObject 是要从中提取数据的 PyTorch Cuda 变量张量，indexes 是要从中检索数据的张量索引。

saveCupy(saveFileName)：将张量保存为 .npy 文件。

loadCupy(loadFileName)：从 .npy 文件加载新的张量。

getNumpyVersion：获取张量的 NumPy 版本。

请参阅此笔记本，了解如何使用数据工具：

https://colab.research.google.com/drive/185Z5Gi62AZxh-EeMfrTtjqxEifHOBXxF

PytorchModelFactory 类

PytorchModelFactory(model_variable,  total_classes,  embed_dimension, datatype = 'float32', deviceType = 'cuda', pinType = False)

使用 PyTorch 张量为模型变量创建切换器。该切换器可以在完整的嵌入集合和模型批次集合之间切换变量。每个变量都需要一个独立的切换器。

示例：

uEmbed_switcher = SpeedTorch.PytorchModelFactory( skip_gram_modelSparse.u_embeddings, total_classes=50000, embed_dimension=128)

参数：

model_variable：您希望为其创建切换器的模型中的特定变量。

total_classes：需要训练的嵌入总数。

embed_dimension：嵌入的维度。

datatype（可选）：变量的数据类型。默认为 'float32'。

deviceType（可选）：设备类型，可设置为 'cuda' 或 'cpu'。默认为 'cuda'。

pinType（可选）：如果设备设置为 'cpu'，可以指定使用固定内存。默认为 False。

方法：

zerosInit()：将变量切换器的完整集合初始化为零。

uniformDistributionInit(low, high)：将变量切换器的完整集合初始化为从 low 到 high 的均匀分布。

normalDistributionInit(mean, stdDev)：将变量切换器的完整集合初始化为均值为 mean、标准差为 stdDev 的正态分布。

customInit(initFunction, *args)：为变量切换器的完整集合使用任何 PyTorch 初始化器。通过 initFunction 传递初始化器，并通过 *args 传递其相应参数。

variableTransformer(batchSize, posPerBatch, negPerBatch = None )：设置用于模型前向传播步骤的占位输入。batchSize 是批次大小，posPerBatch 是每批次的正样本数量。如果还需要用于负样本的第二个占位输入，则可以将 negPerBatch（可选）设置为负样本数量，此时将返回两个占位输入，而不是一个。

beforeForwardPass(retrievedPosIndexes , retrievedNegIndexes = None)：将嵌入从完整嵌入集合切换到模型嵌入集合。retrievedPosIndexes 是要检索的正样本索引。如果也需要检索负样本，则可以同时传递 retrievedNegIndexes（可选）。

afterOptimizerStep(retrievedPosIndexes , retrievedNegIndexes = None)：将更新后的嵌入从模型切换回完整嵌入集合。retrievedPosIndexes 是已检索的正样本索引。如果也检索了负样本，则可以同时传递 retrievedNegIndexes（可选）。

saveTorch(saveFileName)：使用 torch.save 将张量保存到文件。

loadTorch(loadFileName)：使用 torch.load 加载张量。

getNumpyVersion：获取张量的 NumPy 版本。

PytorchOptimizerFactory 类

PytorchOptimizerFactory( given_optimizer,  total_classes,  embed_dimension, model, variable_name, dtype='float32', deviceType = 'cuda', pinType = False)

使用 PyTorch 张量为优化器变量创建切换器。该切换器可以在完整的嵌入集合和优化器批次集合之间切换变量。每个变量都需要一个独立的切换器。

示例：

uAdagrad_switcher = SpeedTorch.PytorchOptimizerFactory(given_optimizer,  total_classes,  embed_dimension, model, variable_name, dtype='float32')

参数：

given_optimizer：使用模型权重初始化的优化器。如果用于嵌入训练，请确保将 sparse 参数设置为 False。目前支持的优化器包括 SparseAdam、Adadelta、Adamax、Adam、AdamW、ASGD 和 RMSprop。Rprop 也被包含在内，但需要先完成第一次前向传播和 loss.backward() 步骤，才能初始化 OptimizerFactory 实例。这是因为 Rprop 优化器需要其参数的梯度来进行初始化。

total_classes：需要训练的嵌入总数。

embed_dimension：嵌入的维度。

model：您的模型实例。

variable_name：模型中定义的变量的确切名称。

dtype（可选）：变量的数据类型。默认为 'float32'。

deviceType（可选）：设备类型，可设置为 'cuda' 或 'cpu'。默认为 'cuda'。

pinType（可选）：如果设备设置为 'cpu'，可以指定使用固定内存。默认为 False。

方法：

optInit：初始化优化器变量切换器。

beforeForwardPass(retrievedPosIndexes , retrievedNegIndexes = None)：将优化器变量权重从完整权重集合切换到优化器权重张量。retrievedPosIndexes 是要检索的正样本索引。如果也需要检索负样本，则可以同时传递 retrievedNegIndexes（可选）。

afterOptimizerStep( retrievedPosIndexes , retrievedNegIndexes = None)：将优化器变量权重从优化器切换回完整权重集合。retrievedPosIndexes 是已检索的正样本索引。如果也检索了负样本，则可以同时传递 retrievedNegIndexes（可选）。

引用 SpeedTorch：

如果您在研究中使用 SpeedTorch 或希望引用，请使用以下格式：

@misc{

title={SpeedTorch},

author={Santosh Gupta},

howpublished={\url{github.com/Santosh-Gupta/SpeedTorch}},

year={2019}

}

SpeedTorch 快速上手指南

SpeedTorch 是一个旨在加速 PyTorch 中 CPU 与 GPU 之间张量传输的开源库。它利用 Cupy 的固定内存（Pinned Memory）技术，在特定场景下（尤其是 CPU 核心数较少时）能显著提升数据加载和稀疏嵌入训练的速度，并支持使用非稀疏优化器进行大规模嵌入训练。

环境准备

在使用 SpeedTorch 之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux (推荐) 或 Windows (需配置 CUDA 环境)。
• Python 版本：建议 Python 3.6+。
• 前置依赖：
  • PyTorch：已安装并可用。
  • CuPy：必须先于 SpeedTorch 安装。CuPy 需要与您的系统 CUDA 版本匹配。
  • GPU 驱动：已正确安装 NVIDIA 驱动和 CUDA Toolkit。

注意：SpeedTorch 的核心加速原理依赖于 CuPy 的索引内核优化。如果您的系统拥有较多 CPU 核心（如 8 核以上），原生 PyTorch 的索引效率可能会提升，此时 SpeedTorch 的相对优势可能减弱，但仍建议在您的具体数据规模下进行基准测试。

安装步骤

请严格按照以下顺序执行安装命令。必须先安装 CuPy，再安装 SpeedTorch。

1. 安装 CuPy

根据您的 CUDA 版本选择对应的 CuPy 安装命令。以下是通用安装方式（自动检测 CUDA 版本）：

pip install cupy-cuda11x
# 如果是 CUDA 12.x，请使用: pip install cupy-cuda12x
# 如果是旧版 CUDA 10.x，请使用: pip install cupy-cuda102

国内用户加速建议：如果下载缓慢，可指定清华或阿里镜像源：

pip install cupy-cuda11x -i https://pypi.tuna.tsinghua.edu.cn/simple

2. 安装 SpeedTorch

pip install SpeedTorch

国内用户加速建议：

pip install SpeedTorch -i https://pypi.tuna.tsinghua.edu.cn/simple

3. 验证安装

在 Python 环境中导入库，确保无报错：

import cupy
import SpeedTorch
print("SpeedTorch 安装成功！")

基本使用

SpeedTorch 的主要用途是加速数据传输以及支持使用常规优化器（如 Adam, Adamax 等）进行稀疏嵌入训练。

场景一：加速 CPU 到 GPU 的数据传输

使用 SpeedTorch 的 DataGadget 可以快速将数据从 CPU 加载到 GPU。

import torch
import cupy
import SpeedTorch

# 初始化 SpeedTorch
# 这一步会设置相关的后端配置
SpeedTorch.init()

# 假设你有一个大型数据集在 CPU 上
# 创建示例数据 (100 万行，128 维)
cpu_data = torch.randn(1000000, 128)

# 使用 SpeedTorch 将数据转换为可在 GPU 高效访问的形式
# 这里演示将数据移至 GPU 并进行操作
gpu_tensor = SpeedTorch.to_cuda(cpu_data)

# 现在 gpu_tensor 可以用于后续的 PyTorch 模型训练
# 当需要传回 CPU 时，速度也会显著快于原生 PyTorch
cpu_result = SpeedTorch.to_cpu(gpu_tensor)

场景二：使用非稀疏优化器训练大规模 Embedding

这是 SpeedTorch 最核心的功能之一。它允许你将巨大的 Embedding 表存放在 CPU 内存中（节省显存），仅在需要时快速传输到 GPU，从而可以使用 Adam, Adamax, RMSprop 等不支持原生稀疏更新的优化器。

import torch
import torch.nn as nn
import SpeedTorch
from SpeedTorch import PytorchModelFactory, PytorchOptimizerFactory

# 1. 定义模型工厂
# 指定设备：'cpu' (使用固定内存以加速传输) 或 'cuda'
# pinned=True 表示启用 CPU 固定内存加速
model_factory = PytorchModelFactory(device='cpu', pinned=True)

# 2. 创建 Embedding 层
# num_embeddings: 词表大小 (可以非常大，受限于 CPU 内存而非显存)
# embedding_dim: 嵌入维度
embedding_layer = model_factory.create_embedding(num_embeddings=10000000, embedding_dim=128)

# 3. 定义优化器工厂
# 现在可以使用任意优化器，例如 Adamax，即使是在做稀疏训练
optimizer_factory = PytorchOptimizerFactory(device='cpu', pinned=True)
optimizer = optimizer_factory.create_optimizer(optimizer_class=torch.optim.Adamax, 
                                               params=embedding_layer.parameters(), 
                                               lr=0.01)

# 4. 训练循环示例
# 假设 input_indices 是当前的批次数据 (位于 CPU 或 GPU 均可，SpeedTorch 会自动处理高效传输)
input_indices = torch.randint(0, 10000000, (32,)) 

# 前向传播 (SpeedTorch 会自动处理必要的 CPU<->GPU 传输)
output = embedding_layer(input_indices)

# 计算损失 (示例)
loss = output.sum()

# 反向传播与更新
loss.backward()
optimizer.step()
optimizer.zero_grad()

print("使用 SpeedTorch + Adamax 完成了一步稀疏训练更新。")

关键提示

• 内存权衡：SpeedTorch 的 CuPy 固定内存张量比原生 PyTorch 张量占用更多内存，但换来了更快的传输速度。
• 适用场景：最适合 CPU 核心数较少（如 2-4 核）且需要处理超大规模 Embedding 表的场景。
• 稀疏标志：在使用上述工厂类创建模型时，由于参数会在每一步都被更新（通过快速传输机制），通常可以将相关参数视为稠密更新处理，从而解锁更多优化器选项。