可以在 CUDA 速度计算中实现 Tensorflow 自定义操作梯度吗？答案

【问题标题】：Can implementing Tensorflow custom op gradients in CUDA speed computation?可以在 CUDA 速度计算中实现 Tensorflow 自定义操作梯度吗？
【发布时间】：2019-03-14 07:46:31
【问题描述】：

在支持 GPU 的 tensorflow 中编写自己的自定义操作时，the guide 建议使用 python 计算梯度。 Elsewhere 人们使用 C++ 和 Eigen 等库以同样的方式更有效地实现渐变。

我的问题是；在 GPU 上训练时，将自定义操作参数张量作为指向设备内存的指针提供（这是正确的吗？）：

OpKernelContext* context
const Tensor& grad = context->input(0);

是否可以通过在 GPU 上计算 win CUDA 操作的梯度来避免主机和设备之间的数据复制？

这会减少计算时间吗？（我知道这取决于梯度计算是否适合并行计算，但假设确实如此）

有什么理由不应该这样做吗？/潜在的速度增加是否太微不足道以至于不值得？

【问题讨论】：

标签： c++ tensorflow eigen

【解决方案1】：

你可以简单地在 CUDA 中构建一个 op，然后在你的梯度的 python 定义中调用它，这样你肯定可以大大加快梯度计算，而不必在 GPU 和 CPU 内存之间复制

【讨论】：

在 CUDA 中构建操作需要您明确指定要由 GPU 操作的数据的内存位置。在所有示例中，此数据源自主机，您需要将其复制到 gpu，然后对其进行操作，然后将其复制回主机以访问它。但是，如果数据之前已经使用 CUDA 在 tensorflow 中进行了计算，那么应该可以访问 gpu 上已经存在的数据并将其留在那里以供将来的操作访问，从而避免昂贵的数据传输。
我不明白为什么你的输入现在在主机上，你不是说它们在 GPU 上吗？您可以在 GPU 上完全训练模型：向前和向后传递。