How to simulate reduced precision floats in TensorFlow?

Question

I would like a way to reduce the precision of floats in TensorFlow (approximately: truncate the mantissa) to an arbitrary number of bits within a defined full range. I don't need to write code entirely in reduced precision (like tf.float16), but rather to come up with a series of operations that reduce the precision of a tensor while leaving it the original type (eg tf.float32).

For example, if the full range is 0 to 1, and the precision is 8 bit, 0.1234 would become round(0.1234 * 256) / 256 = 0.125. This uses simple rounding.

I would also like to do statistical rounding, where the probability of rounding in each direction is proportional to how far the value is from that. For example, 0.1234 * 256 = 31.5904, which would round up to 32/256 59% of the time, and to 31/256 41% of the time.

Extra question: How to take an existing graph and modify it to add rounding after every convolution?

Answer 1

The only tricky part is to provide the gradients to the rounding operation. The already implemented tf.round does not have a gradient implemented. But you can implement your own rounding operation (statistical or simple rounding both work) as shown here: Tensorflow: How to write op with gradient in python?

Where you can simply use:

grad(round(T)) = round(grad(T))

Now once you have your personalized round operation which transfers gradients you can simply do:

def reduce_precision(tensor, precision_bits=8):
    N = 2**precision_bits
    return round(N * tensor)/N

And for the stochastic rounding, you can create a simple numpy function like

def stochastic_round(x):
    r,f = np.modf(x)
    return r + np.random.binomial(1,r)

and then tensoflow-ize it as shown in How to make a custom activation function with only Python in Tensorflow?

where you can define it's gradient operation as

def grad_stochastic_round(op, grad):
    return stochastic_round(grad)