基于KL散度的INT8训练后量化

我们知道，P,Q两列数据的相对熵越小，那么P,Q分布越接近，用Q近似P损失的信息就少，英伟达的INT8量化就是基于这个原理，如图是英伟达int8量化的算法伪代码

下面是根据相对熵来选取最佳阈值的代码。

import numpy as np
import copy

def compute_kl_divergence(P,Q):
    length=len(P)
    sum=0.0
    for i in range(length):
        if P[i]!=0:
            if Q[i]==0:
                sum+=1
            else:
                sum+=P[i]*np.log(P[i]/Q[i])
    return sum


def threshold_distribution(distribution,target_bin):
    target_threshold = target_bin
    min_kl_divergence = 10000000000000
    length = len(distribution)



    for threshold in range(target_bin,length):
        #t_distribution=np.empty((threshold,))
        t_distribution=copy.deepcopy(distribution[0:threshold])
        t_distribution[threshold - 1] += np.sum(distribution[threshold:])

        #get P
        num_per_bin = threshold / target_bin

        quantize_distribution = np.zeros((target_bin,))

        for i in range(target_bin):
            start = i * num_per_bin
            end = start + num_per_bin

            left_upper = int(np.ceil(start))
            if left_upper > start:
                left_scale = left_upper - start
                quantize_distribution[i] += left_scale * distribution[left_upper - 1]
            right_lower = int(np.floor(end))

            if right_lower < end:
                right_scale = end - right_lower
                quantize_distribution[i] += right_scale * distribution[right_lower]

            for j in range(left_upper,right_lower):
                quantize_distribution[i] += distribution[j]

        # get Q
        expand_distribution=np.zeros_like(t_distribution)

        for i in range(target_bin):
            start = i * num_per_bin
            end = start + num_per_bin

            count = 0

            left_upper = int(np.ceil(start))
            left_scale = 0
            if left_upper > start:
                left_scale = left_upper - start
                if t_distribution[left_upper - 1] != 0:
                    count += left_scale

            right_lower = int(np.floor(end))
            right_scale = 0
            if right_lower < end:
                right_scale = end - right_lower
                if t_distribution[right_lower] != 0:
                    count += right_scale

            for j in range(left_upper,right_lower):
                if t_distribution[j] != 0:
                    count+=1

            expand_value = quantize_distribution[i] / count

            if left_upper > start:
                if t_distribution[left_upper - 1] != 0:
                    expand_distribution[left_upper - 1] += expand_value * left_scale
            if right_lower < end:
                if t_distribution[right_lower] != 0:
                    expand_distribution[right_lower] += expand_value * right_scale
            for j in range(left_upper,right_lower):
                if t_distribution[j] != 0:
                    expand_distribution[j] += expand_value

        kl_divergence = compute_kl_divergence(t_distribution, expand_distribution)

        #print(threshold,kl_divergence)

        if kl_divergence < min_kl_divergence:
            min_kl_divergence = kl_divergence
            target_threshold = threshold

    return target_threshold


if __name__=='__main__':
    distribution=np.empty((2048,))
    for i in range(len(distribution)):
        distribution[i]=i
    distribution/=np.sum(distribution)
    target_threshold=threshold_distribution(distribution,128)
    print(target_threshold)