[PTQ]均匀量化和非均匀量化

news2024/11/16 13:06:51

均匀量化和非均匀量化

基本概念

量化出发点：使用整型数据类型代替浮点数据，从而节省存储空间同时加快推理速度。
量化基本形式
- 均匀量化：浮点线性映射到定点整型上，可以根据scale/offset完成量化/反量化操作。
- 非均匀量化
  - PowersOfTwoQuant：浮点映射成2的指数位上，根据power计算完成量化/反量化操作。
  - AdditivePowersOfTwoQuant：浮点映射成多个不同指数的和，而整型表示应该是跟硬件实现相关（后续待定聊聊）
不同形式的量化对不同的数据分布的拟合程度不一样

均匀量化(UniformQuant)

特点：
- 适合range较大，分布较均匀的分布
- 最普遍/大众的量化方式，基本上市面上的加速器都是只支持均匀量化
- 可以很容易从对称量化扩展到非对称量化

python代码实现如下所示

# 均匀量化
class UniformQuant:
    def __init__(self, 
                bit=8) -> None:
        self.bit = bit

    def quant(self, input, min=0.0, max=1.0):
        max_q = np.power(2, self.bit) - 1
        input = np.clip(input, min, max)
        scale = (max - min) / max_q
        q = np.clip(np.round(input / scale), 
                    0,
                    max_q)
        return q, scale

    def dequant(self, q, scale):
        return (scale * q).astype(np.float32)

    def fakequant(self, input):
        q, scale = self.quant(input)
        fakequant_input = self.dequant(q, scale)
        return fakequant_input

PowersOfTwoQuant

特点：
- 适合以0.0为中心的高斯分布数据，靠近0.0附近的精度非常高
- 提高bits数，无法提高整体的精度

python代码实现如下所示

# PoT量化
class PowersOfTwoQuant:
    def __init__(self, 
                bit=8) -> None:
        self.bit = bit

    def quant(self, input, alpha):
        max_q = np.power(2, self.bit) - 1
        input = np.clip(input / alpha, 0.0, 1.0)
        q = list()
        for i in input:
            if i == 0.0:
                q.append(0)
            else:
                e_tmp = np.log(i) / np.log(2.0)
                e_tmp = np.round(e_tmp + max_q)
                e_tmp = np.clip(e_tmp, 0, max_q)
                q.append(e_tmp)
        q = np.array(q, dtype=np.uint32)
        return q

    def dequant(self, q, alpha):
        max_q = np.power(2, self.bit) - 1
        input = list()
        for i in q:
            if i == 0:
                input.append(0.0)
            else:
                input.append(np.power(2.0, i - max_q))
        
        input = np.array(input, dtype=np.float32)
        input = alpha * input
        return input

    def fakequant(self, input, alpha = None):
        if alpha is None:
            alpha = input.max()
        q = self.quant(input, alpha)
        fakequant_input = self.dequant(q, alpha)
        return fakequant_input

AdditivePowersOfTwoQuant

特点：
- 0.0附近的精度比均匀分布高
- 提高bits数，可以提高整体的精度
- 计算较为复杂，推理加速应该需要专门设计实现方式

python代码实现如下所示

# APoT量化
class AdditivePowersOfTwoQuant:
    def __init__(self, 
                bit=8,
                k=2) -> None:
        self.bit = bit
        self.k = k
        self.n = int(self.bit/self.k)
        self.p = np.array(self.gen_table(), dtype=np.float32)
        self.sort_values = self.get_all_values()
        self.sort_values.sort()
        # self.power_set = build_power_value(self.bit, additive=False)
        # self.power_set = self.power_set.detach().numpy()

    def get_all_values(self):
        m,n = list(self.p.shape)
        values = self.recursion_add(self.p.tolist())
        return np.array(values, dtype=np.float32)
    

    def recursion_add(self, input_list, callback=lambda a,b:a+b):
        # [list1, list2]
        res = list()
        if len(input_list) == 2:
            list1, list2 = input_list
            for a in list1:
                for b in list2:
                    res.append(callback(a,b))
            return res
        elif len(input_list) > 3:
            tmp_input_list = input_list[:-2]
            tmp_res = self.recursion_add(input_list[-2:])
            tmp_input_list.append(tmp_res)
            return self.recursion_add(tmp_input_list)
        

    def gen_table(self):
        p = list()
        for i in range(self.n):
            one_lines_p = list()
            one_lines_p.append(0.0)
            for j in range(int(math.pow(2, self.k) - 2 + 1)):
                one_lines_p.append(math.pow(2, -(i+j*self.n)))
            p.append(one_lines_p)
        return p

    def quant(self):
        # 待定
        pass

    def dequant(self):
        # 待定
        pass

    def fakequant(self, input):
        scale = input.max() / self.sort_values.max()
        fakequant_input = list()
        input = input / scale
        for v in input:
            argmin_index = np.argmin(np.abs(self.sort_values - v))
            tmp_value = self.sort_values[argmin_index]
            fakequant_input.append(tmp_value)
        fakequant_input = np.array(fakequant_input, dtype=np.float32) * scale
        return fakequant_input