HWGQ-TF

HWGQ-TF.py

def get_hwgq(bitA):

    def quantize(x, k):
        # in order of 
        assert k in [2,3,4,5], 'Does not support %d bits' % k
        code_book={
        '2':[0.5380, 0., 0.5380*(2**2-1)],
        '3':[0.3218, 0., 0.3218*(2**3-1)],
        '4':[0.1813, 0., 0.1813*(2**4-1)],
        '5':[0.1029, 0., 0.1029*(2**5-1)]
        }
        delta, minv, maxv = code_book[str(k)]
        #print(delta,minv,maxv)
        @tf.custom_gradient
        def _quantize(x):
            return tf.to_float(x>0.)*(tf.clip_by_value((tf.floor(x/delta + 0.5)+tf.to_float(x<0.5*delta))*delta, minv, maxv)), lambda dy: dy*tf.to_float(x>minv)*tf.to_float(x<maxv)

        return _quantize(x)

    def fa(x):
        if bitA == 32:
            return x

        return quantize(x, bitA)
    return fa