TensorFlow pass 梯度不变

Question

假设我在神经网络中使用了一些自定义操作binarizer。该操作采用Tensor 并构造一个新的Tensor。我想修改该操作，使其仅用于前向传递。在反向传播中，当计算梯度时，它应该只是通过到达它的梯度。

更具体的说，binarizer 是：

def binarizer(input):
    prob = tf.truediv(tf.add(1.0, input), 2.0)
    bernoulli = tf.contrib.distributions.Bernoulli(p=prob, dtype=tf.float32)
    return 2 * bernoulli.sample() - 1

然后我设置了我的网络：

# ...

h1_before_my_op = tf.nn.tanh(tf.matmul(x, W) + bias_h1)
h1 = binarizer(h1_before_b)

# ...

loss = tf.reduce_mean(tf.square(y - y_true))
train_step = tf.train.GradientDescentOptimizer(0.5).minimize(loss)

如何告诉 TensorFlow 在反向传递中跳过梯度计算？

---

我尝试按照this answer 中的描述定义自定义操作，但是：py_func 不能返回 Tensors，这不是它的用途——我明白了：

UnimplementedError（回溯见上文）：不支持的对象类型张量

Answer 1

你正在寻找tf.stop_gradient(input, name=None):

停止梯度计算。

当在图中执行时，此操作按原样输出其输入张量。

h1 = binarizer(h1_before_b)
h1 = tf.stop_gradient(h1)