How to backprop for custom quantization in tensorflow 2.0?

user_13127499 0
Viewed 417

I would like to backprop with identity (straight through estimator) when I quantize a tensor in tf2.0 eager mode. I am able to perform the forward pass, but encounter ValueError: No gradients provided for any variable.. error when trying to apply gradients for my model for backprop.

I tried to use custom gradients. As a minimal example, consider the following code:

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

The gradient is None, which means I won't be able to backprop when I use this kind of functions in a model. How should I fix this? Are there better ways to achieve quantization given that I need something other than tf.quantize()?

Also, the way I used the GradientTape g is according to the official tutorials, but I don't understand why they would use g outside the with scope (last line in the code). Any clarification is appreciated.

quantization tensorflow2.0 tensorflow machine-learning python
1 Answers

The problem is that you quantize by converting to integer. Use the floor function instead

@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)
user_105651970
Related