基于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  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  start:
                if t_distribution[left_upper - 1] != 0:
                    expand_distribution[left_upper - 1] += expand_value * left_scale
            if right_lower