pytorch框架下的参数量化（量化为任意位数，或2的幂次方）

pytorch框架下参数渐进量化的实现

将pytorch框架下的参数量化为特定形式，会产生一定的误差，这篇博客以MINIST数据集，LSTM量化为例，主要写了量化的详细流程，并附上完整程序。

---

一、量化原理

本博客介绍的量化方式，可以将参数量化成任何形式，但量化后的参数不支持反向传播，即不能再训练。

LSTM共有8个权重矩阵，如果直接量化，会产生较大的误差，因此本博客采用渐进量化的策略来减少误差。量化流程如下：

• 1. 先训练模型，等待收敛
• 2. 量化Wi，将量化后的Wi保存为常数，对模型进行再训练（此时模型可训练参数只有7个权重矩阵和4个bias）
• 3. 量化Ui，将量化后的Ui保存为常数，对模型进行再训练（此时模型可训练参数只有6个权重矩阵和4个bias）
• 4. 依此类推，最终量化引起的误差，只有量化最后一个权重矩阵产生的误差

以下是完整的程序实现流程

---

二、自定义RNN框架

pytorch自带的nn.LSTM只有两个权重矩阵，他将Wi，Wf，Wo，Wc四个矩阵合并成一个矩阵，将Ui，Uf，Uo，Uc四个矩阵合并成一个矩阵，因此不适合这种量化方式，需要自定义LSTM层；

• 自定义RNN代码：自定义RNN结构
  自定义LSTM部分代码如下：

class LSTMCell(RNNCellBase):

    def __init__(self, input_size, hidden_size, bias=True, grad_clip=None):
        super(LSTMCell, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.grad_clip = grad_clip
        self.weight_ix = Parameter(torch.Tensor(hidden_size, input_size))
        self.weight_ih = Parameter(torch.Tensor(hidden_size, hidden_size))
        self.weight_cx = Parameter(torch.Tensor(hidden_size, input_size))
        self.weight_ch = Parameter(torch.Tensor(hidden_size, hidden_size))
        self.weight_ox = Parameter(torch.Tensor(hidden_size, input_size))
        self.weight_oh = Parameter(torch.Tensor(hidden_size, hidden_size))
        self.weight_fx = Parameter(torch.Tensor(hidden_size, input_size))
        self.weight_fh = Parameter(torch.Tensor(hidden_size, hidden_size))
        if bias:
            self.bias_i = Parameter(torch.Tensor(hidden_size))
            self.bias_f = Parameter(torch.Tensor(hidden_size))
            self.bias_g = Parameter(torch.Tensor(hidden_size))
            self.bias_o = Parameter(torch.Tensor(hidden_size))
        else:
            self.register_parameter('bias', None)

        self.reset_parameters()

    def reset_parameters(self):
        stdv = 1.0 / math.sqrt(self.hidden_size)
        for weight in self.parameters():
            weight.data.uniform_(-stdv, stdv)

    def forward(self, input, hx):
        h, c = hx
        i = F.linear(input, self.weight_ih, self.bias_i) + F.linear(h, self.weight_ih)
        f = F.linear(input, self.weight_fh, self.bias_f) + F.linear(h, self.weight_fh)
        g = F.linear(input, self.weight_gh, self.bias_g) + F.linear(h, self.weight_gh)
        o = F.linear(input, self.weight_oh, self.bias_o) + F.linear(h, self.weight_oh)

        if self.grad_clip:
            i = clip_grad(i, -self.grad_clip, self.grad_clip)
            f = clip_grad(f, -self.grad_clip, self.grad_clip)
            g = clip_grad(g, -self.grad_clip, self.grad_clip)
            o = clip_grad(o, -self.grad_clip, self.grad_clip)
        i = F.sigmoid(i)
        f = F.sigmoid(f)
        g = F.tanh(g)
        o = F.sigmoid(o)

        c = f * c + i * g
        h = o * F.tanh(c)
        return h, c

含有8个权重矩阵和bias的可训练参数。

---

三、MNIST数据集和建模，初始化

这篇博客采用的模型如下：

class RNNModel(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, num_classes, bias=True, grad_clip=None):
        super(RNNModel, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers

        self.rnn = LSTM(input_size, hidden_size, num_layers=num_layers, 
                        bias=bias, return_sequences=False, grad_clip=grad_clip)
        self.fc = nn.Linear(hidden_size, num_classes, bias=bias)
    
    def forward(self, x):
        # Set initial states 

        zeros = Variable(torch.zeros(x.size(0), self.hidden_size))
        initial_states = [(zeros, zeros)] * self.num_layers
        
        # Forward propagate RNN
        out, _ = self.rnn(x, initial_states)  
        # Decode hidden state of last time step
        out = self.fc(out)

        return out

model = RNNModel(input_size, hidden_size, num_layers, num_classes, bias=True, grad_clip=10)

model的类型如图所示：

检查一下model中的可训练参数：

for name,param in model.named_parameters():
    print(name)  #输出所有可训练惨的名称

运行结果如图所示：

我们量化的目标就是前8个参数

---

四、量化函数介绍

这篇博客提供两种量化函数：
1，t1 将函数量化为m位整数，f位小数的特定位数。
2，t2 将参数量化为2的幂次方。
两种量化函数的详解参考这篇博客：量化函数

def t1(a):             #输入训练参数，输出量化后的训练参数
    b = a.detach().numpy() #由于a是训练参数，requires_grad为True，因此不能直接用numpy函数操作，需转换
    b = np.clip(b,-0.9995117187,0.9995117187) #0.9995117187是1 - (1/2)^11  
    b = np.round(b * 2048 + 0.5) / 2048    #2048是2^11
    a.data = torch.from_numpy(b).data     #得到最接近原始a的定点数
    return a


def t2(a):
    b = a.detach().numpy()         
    e = np.sign(b)
    b = np.clip(np.round(np.log2(np.fabs(b))+0.4),-7,0) #得到最接近原始a的2的幂次方，不改变a的其他属性，因此只使用data属性
    b = np.power(2,b) * e
    a.data = torch.from_numpy(b).data
    return a
#t1 12位定点量化，整数位数0，小数位数11，符号位1位
#t2 量化为2的幂次方，整数位数0，小数位数7，符号位1位

---

五、量化权重矩阵

首先将未量化的模型训练至收敛；

state = {'net': model.state_dict()} #将模型保存为字典变量
del state["net"]["rnn.cell0.weight_ih"]#将字典里的weight_ih权重删掉
torch.save(state,"save/quantization_ih.pt")#保存模型，这里模型里没有weight_ih函数了
weight_ih = {"model.rnn.cell0.weight_ih":t2(model.rnn.cell0.weight_ih)}#将weight_ih量化为2的幂次方形式，并保存
torch.save(state,"save/weight_ih.pt")

此时我们获得两个保存的文件quantization_ih.pt文件不包含weight_ih矩阵了，weight_ih.pt中保存的是量化为2的幂次方后的weight_ih矩阵；下一步我们要将自定义LSTM中的weight_ih矩阵变成不可训练的参数；

'''self.weight_ih = Parameter(torch.Tensor(hidden_size, hidden_size))'''
self.weight_ih = torch.Tensor(hidden_size, hidden_size)#删掉自定义LSTM文件里的Parameter(),这样weight_ih就只是模型的一个变量，而不是可训练参数了

check = torch.load("save/quantization_ih.pt")#加载不含weight_ih的模型
model.load_state_dict(check["net"])
model.rnn.cell0.weight_ih = torch.load("save/weight_ih.pt")#加载量化后的weight_ih矩阵

之后继续训练模型，此时的weight_ih就是固定的了，可训练参数只有其他权重矩阵和bias了，将训练好的模型保存；

state = {'net': model.state_dict()} #将模型保存为字典变量
del state["net"]["rnn.cell0.weight_oh"]#将字典里的weight_oh权重删掉
torch.save(state,"save/quantization_ih_oh.pt")#保存模型，这里模型里没有weight_ih和weight_oh矩阵了
weight_oh = {"model.rnn.cell0.weight_oh":t2(model.rnn.cell0.weight_oh)}#将weight_oh量化为2的幂次方形式，并保存
torch.save(state,"save/weight_oh.pt")

'''self.weight_oh = Parameter(torch.Tensor(hidden_size, hidden_size))'''
self.weight_oh = torch.Tensor(hidden_size, hidden_size)#再删掉自定义LSTM文件里的Parameter(),这样weight_ih和weight_oh就只是模型的一个变量，而不是可训练参数了

之后继续训练模型，此时的weight_oh和weight_ih就是固定的了，可训练参数只有其他权重矩阵和bias了，将训练好的模型保存；

check = torch.load("save/quantization_ih_oh.pt")#加载不含weight_ih和weight_oh训练好的模型
model.load_state_dict(check["net"])
model.rnn.cell0.weight_ih = torch.load("save/weight_ih.pt")#加载量化后的weight_ih矩阵
model.rnn.cell0.weight_oh = torch.load("save/weight_oh.pt")#加载量化后的weight_oh矩阵

依次类推，逐渐完成所有权重的量化

总结

这篇主要介绍了一种pytorch框架下的权重参数量化方式，有点绕，手动画个流程图，仅供供参考：

最后再捋一捋思路，

• 第一步，训练，模型
• 第二步，删除模型的Wih的矩阵，然后保存模型，并将Wih矩阵量化后单独保存
• 第三步，删除自定义RNN中Wih的Parameter属性（此时Wih变成常数，不会再被训练了）
• 第四步，加载训练好的，没有Wih的模型，并将单独保存的量化后的Wih赋值给模型的Wih
• 第五步，训练模型，此时Wih是被量化后的数值，而且不会变了（不参与梯度下降）
• 第六步，在此基础上，量化第二个矩阵
• 以此类推，量化完全部的矩阵

示例工程代码：

• pytorch-LSTM-quantization