TensorFlow含gather操作的函数梯度：IndexedSlices转Tensor问题

这问题我之前踩过坑！当计算涉及tf.gather这类索引操作的梯度时，TensorFlow会返回IndexedSlices对象而非常规张量，本质是为了高效存储稀疏梯度——毕竟gather只用到了参数的部分索引，其他位置梯度都是0，用稀疏表示能省内存。下面给你几种靠谱的解决办法：

方法1：用官方公开的tf.convert_to_tensor()（推荐） #

这是最稳妥的方式，TensorFlow的公开API已经内置了对IndexedSlices的转换支持，直接把梯度转成稠密张量就行。修改你的梯度计算代码如下：

# 先获取梯度（注意tf.gradients返回的是列表，取第一个元素）
raw_grad = tf.gradients(T_loss, [T_W])[0]
# 转换为稠密张量
T_grad = tf.convert_to_tensor(raw_grad)

方法2：用内部的_IndexedSlicesToTensor（不推荐长期用） #

你注意到的tensorflow.python.ops.gradients_impl._IndexedSlicesToTensor确实能转换，但它是下划线开头的内部函数，后续TensorFlow版本可能会修改或移除，所以只适合临时调试用：

from tensorflow.python.ops import gradients_impl as GI
raw_grad = tf.gradients(T_loss, [T_W])[0]
T_grad = GI._IndexedSlicesToTensor(raw_grad)

修改后的完整示例代码 #

把你的代码按方法1修改后，运行就能得到预期的2元素梯度向量：

import tensorflow as tf
import numpy as np

T_W = tf.placeholder(tf.float32, [2], 'W') # 参数向量
T_data = tf.placeholder(tf.float32, [10], 'data') # 数据向量
T_Di = tf.placeholder(tf.int32, [10], 'Di') # 索引向量
T_pred = tf.gather(T_W, T_Di)
T_loss = tf.reduce_sum(tf.square(T_data - T_pred)) # 损失函数

# 转换梯度为稠密张量
raw_grad = tf.gradients(T_loss, [T_W])[0]
T_grad = tf.convert_to_tensor(raw_grad)

init = tf.global_variables_initializer()
with tf.Session() as sess:
    sess.run(init)
    feed_dict = {
        T_W: [1., 2.], 
        T_data: np.arange(10)**2, 
        T_Di: np.arange(10)%2
    }
    loss_val, grad_val = sess.run([T_loss, T_grad], feed_dict=feed_dict)
    print(f"损失值: {loss_val}")
    print(f"梯度向量: {grad_val}")

运行后会输出类似这样的结果（和手动计算的梯度一致）：

损失值: 28730.0
梯度向量: [-340. -404.]

补充说明 #

为什么会返回IndexedSlices？简单说：tf.gather的梯度是稀疏的——只有被索引到的参数位置有非零梯度，其他位置都是0。TensorFlow用IndexedSlices存储三个信息：非零梯度值、对应的索引、原参数的形状，比直接存储全零的稠密张量更高效，尤其当参数维度很大时。但如果你的参数维度很小（比如示例里的2维），转成稠密张量完全没性能问题。

内容的提问来源于stack exchange，提问作者segasai