解决SparseApply *的基于gpu-gradient的稀疏更新的变通方法 - 张量流中有没有？

Question

从一些背景开始：

• SparseApply*运算符未针对GPU实现（从0.11开始）
• 使用Momentum和ADAM以及各种其他基于渐变的优化器
需要这些操作
• CPU版本导致CPU和GPU之间的数据传输和延迟

我希望Momentum / ADAM在一个在许多地方广泛使用tf.gather的模型上，其方式不会削弱GPU利用率。我的变量的密集版本并不是特别大，特别是与计算图中的激活和其他张量使用的内存相比。

如果没有为GPU实现这些SparseApply*操作系统，我想到了一个解决方法，我想对它是否有可能起作用，替代建议或只是一般性批评进行一些评论：

1. 在dict D
中跟踪模型中的所有模型参数变量
2. 使用D
收集tf.gradients中参数的渐变
3. 使用子图H中的密集GPU操作（及其潜在变量）手动编程adam /动量更新。（整个张量流/会话默认计算图的子图）
4. 计算/应用H

一些背景信息：我有选择地从https://github.com/tensorflow/tensorflow/issues/2314的stephenroller的第一条评论中选择了一个子字符串

  

目前，似乎只有GradientDescentOptimizer支持   当存在SparseTensor更新时在GPU上运行。

     

允许它们可以大幅度提高速度   存储在GPU上。首先，不需要转移载体   嵌入到GPU中或从GPU嵌入，而是可以只转移   嵌入索引，然后渐变也不需要   向后转移。在一次测试中，我运行了我实现的地方   在GPU上的版本，一个时代的差异大约是1100秒vs   300秒。

编辑：重命名我的问题/标题，以便在标题的前面显示关键字。

Answer 1

另一个选择是使用普通的tensorflow操作来实现这些优化器，并使用scatter_add进行最后的添加，这应该可以工作，因为scatter_add已在GPU中注册。