python - 如何在 tensorflow 2.0 中为自定义量化反向传播？

Question

当我在 tf2.0 eager 模式下量化张量时，我想使用恒等式进行反向传播(直接通过估计器)。我能够执行正向传递，但遇到 ValueError: No gradients provided for any variable.. 尝试为我的反向传播模型应用梯度时出现错误。

我尝试使用自定义渐变。作为最小示例，请考虑以下代码:

import tensorflow as tf
@tf.custom_gradient
def quantize(x):
    x = tf.cast(x, dtype=tf.uint8)
    def grad(dy):
        return dy
    return x, grad

@tf.custom_gradient
def dequantize(x):
  x = tf.cast(x, dtype=tf.float32)
  def grad(dy):
       return dy
  return x, grad

x = tf.ones([10,10])
with tf.GradientTape() as g:
  g.watch(x)
  y = dequantize(quantize(x))*2
dy_dx = g.gradient(y, x) 

print(dy_dx) # outputs None

梯度为 None，这意味着当我在模型中使用这种函数时，我将无法反向传播。我应该如何解决这个问题？如果我需要 tf.quantize() 以外的东西，是否有更好的方法来实现量化？

另外，我使用GradientTape g的方法是按照官方教程，但是我不明白为什么他们会在with之外使用g 范围(最后代码中的行)。任何澄清表示赞赏。

Answer 1

问题是您通过转换为整数来量化。改为使用 floor 函数

@tf.custom_gradient
def quantize(x):
    x = tf.math.floor(x)
    def grad(dy):
        return dy
    return x, grad

x = tf.ones((10,10))
with tf.GradientTape() as g:
  g.watch(x)
  y = quantize(1.1*x+.1)*2

dy_dx = g.gradient(y, x, output_gradients = tf.ones((10,10)))

print(dy_dx)