Return the three quantization functions fw, fa, fg, for weights, activations and gradients respectively
(bitW, bitA, bitG)
| 6 | |
| 7 | |
| 8 | def get_dorefa(bitW, bitA, bitG): |
| 9 | """ |
| 10 | Return the three quantization functions fw, fa, fg, for weights, activations and gradients respectively |
| 11 | """ |
| 12 | def quantize(x, k): |
| 13 | n = float(2 ** k - 1) |
| 14 | |
| 15 | @tf.custom_gradient |
| 16 | def _quantize(x): |
| 17 | return tf.round(x * n) / n, lambda dy: dy |
| 18 | |
| 19 | return _quantize(x) |
| 20 | |
| 21 | def fw(x): |
| 22 | if bitW == 32: |
| 23 | return x |
| 24 | |
| 25 | if bitW == 1: # BWN |
| 26 | E = tf.stop_gradient(tf.reduce_mean(tf.abs(x))) |
| 27 | |
| 28 | @tf.custom_gradient |
| 29 | def _sign(x): |
| 30 | return tf.where(tf.equal(x, 0), tf.ones_like(x), tf.sign(x / E)) * E, lambda dy: dy |
| 31 | |
| 32 | return _sign(x) |
| 33 | |
| 34 | x = tf.tanh(x) |
| 35 | x = x / tf.reduce_max(tf.abs(x)) * 0.5 + 0.5 |
| 36 | return 2 * quantize(x, bitW) - 1 |
| 37 | |
| 38 | def fa(x): |
| 39 | if bitA == 32: |
| 40 | return x |
| 41 | return quantize(x, bitA) |
| 42 | |
| 43 | def fg(x): |
| 44 | if bitG == 32: |
| 45 | return x |
| 46 | |
| 47 | @tf.custom_gradient |
| 48 | def _identity(input): |
| 49 | def grad_fg(x): |
| 50 | rank = x.get_shape().ndims |
| 51 | assert rank is not None |
| 52 | maxx = tf.reduce_max(tf.abs(x), list(range(1, rank)), keep_dims=True) |
| 53 | x = x / maxx |
| 54 | n = float(2**bitG - 1) |
| 55 | x = x * 0.5 + 0.5 + tf.random_uniform( |
| 56 | tf.shape(x), minval=-0.5 / n, maxval=0.5 / n) |
| 57 | x = tf.clip_by_value(x, 0.0, 1.0) |
| 58 | x = quantize(x, bitG) - 0.5 |
| 59 | return x * maxx * 2 |
| 60 | |
| 61 | return input, grad_fg |
| 62 | |
| 63 | return _identity(x) |
| 64 | return fw, fa, fg |
| 65 |
no outgoing calls
no test coverage detected
searching dependent graphs…