深度学习模型压缩#2-量化

深度学习模型压缩

Qinghuan1958

795人浏览 · 2022-05-27 16:26:52

Qinghuan1958 · 2022-05-27 16:26:52 发布

深度学习模型压缩

上接
模型的剪枝是为了减少参数量和运算量，而量化是为了压缩数据的占用量。

量化

所谓的模型量化就是将浮点存储（运算）转换为整型存储（运算）的一种模型压缩技术。
好处：提升计算效率；减少内存和存储占用；减少能耗。

对于量化对象，主要有以下三个，实际中可能是量化其中的多个甚至全部：
weight（权重）：weight的量化是最常规也是最常见的。量化weight可达到减少模型大小和memory footprint等目的。
activation（激活函数输出）：实际中activation往往是占内存使用的大头，因此量化activation不仅可以大大减少memory footprint。更重要的是，结合weight的量化可以充分利用整数计算获得性能提升。
gradient（梯度）：相对上面两者略微小众一些，因为主要用于训练。它主要作用是在分布式计算中减少通信开销，单机训练时也可减少backward时的开销。

对于量化位数，可以有很多种选择。大体可分为几类：
float16量化是比较保险的做法，大多数情况下可以不丢失太多精度的情况下有明显的性能提升。因为都是浮点，相对来说容易，就不细展开了。
8位是比较常见的，也是相对成熟的。相关的研究很多，各种主流框架也基本都支持。
8位以下目前而言学界相对玩得多些，工业界有少量支持，但还没有太成熟。8位以下主要是4，2和1位（因为位数为2的幂次性能会更好，也更容易实现）。如果精度低至1位，也就是二值化，那可以用位运算进行计算。这对处理器而言是很友好的。

模型量化实现步骤
1-在输入数据（通常是权重或者激活值）中统计出相应的min_value和max_value;
2-选择合适的量化类型，对称量化（int8）还是非对称量化（uint8）；
3-根据量化类型、min_value和max_value来计算获得量化的参数Z/Zero point和S/Scale；
4-根据标定数据对模型执行量化操作，即将其由FP32转换为INT8；
5-验证量化后的模型性能，如果效果不好，尝试着使用不同的方式计算S和Z,重新执行上面的操作；

Pytorch模型量化

# 导入第三方的库函数
import os
from io import open
import time

import torch
import torch.nn as nn
import torch.quantization
import torch.nn.functional as F

# 创建LSTM模型类
class LSTMModel(nn.Module):
    """整个网络包含一个encoder, 一个recurrent模块和一个decoder."""

    def __init__(self, ntoken, ninp, nhid, nlayers, dropout=0.5):
        super(LSTMModel, self).__init__()
        # 预定义一些网络层
        self.drop = nn.Dropout(dropout)
        self.encoder = nn.Embedding(ntoken, ninp)
        self.rnn = nn.LSTM(ninp, nhid, nlayers, dropout=dropout)
        self.decoder = nn.Linear(nhid, ntoken)
        self.init_weights()
        self.nhid = nhid
        self.nlayers = nlayers

    def init_weights(self):
    	'''
    	初始化模型权重
    	'''
        initrange = 0.1
        self.encoder.weight.data.uniform_(-initrange, initrange)
        self.decoder.bias.data.zero_()
        self.decoder.weight.data.uniform_(-initrange, initrange)

    def forward(self, input, hidden):
    	'''
    	搭建网络并执行前向推理
    	'''
        emb = self.drop(self.encoder(input))
        output, hidden = self.rnn(emb, hidden)
        output = self.drop(output)
        decoded = self.decoder(output)
        return decoded, hidden

    def init_hidden(self, bsz):
    	'''
    	初始化hidden层的权重
    	'''
        weight = next(self.parameters())
        return (weight.new_zeros(self.nlayers, bsz, self.nhid),
                weight.new_zeros(self.nlayers, bsz, self.nhid))

# 创建一个词典类，用来处理数据
class Dictionary(object):
    def __init__(self):
        self.word2idx = {}
        self.idx2word = []

    def add_word(self, word):
    	'''
    	在词典中添加新的word
    	'''
        if word not in self.word2idx:
            self.idx2word.append(word)
            self.word2idx[word] = len(self.idx2word) - 1
        return self.word2idx[word]

    def __len__(self):
    	'''
    	返回词典的长度
    	'''
        return len(self.idx2word)

class Corpus(object):
    def __init__(self, path):
        self.dictionary = Dictionary()
        # 分别获取训练集、验证集和测试集
        self.train = self.tokenize(os.path.join(path, 'train.txt'))
        self.valid = self.tokenize(os.path.join(path, 'valid.txt'))
        self.test = self.tokenize(os.path.join(path, 'test.txt'))

    def tokenize(self, path):
        """对输入的文件执行分词操作"""
        assert os.path.exists(path)
        # 将新的单词添加到词典中
        with open(path, 'r', encoding="utf8") as f:
            for line in f:
                words = line.split() + ['<eos>']
                for word in words:
                    self.dictionary.add_word(word)

        # 标记文件的内容
        with open(path, 'r', encoding="utf8") as f:
            idss = []
            for line in f:
                words = line.split() + ['<eos>']
                ids = []
                for word in words:
                    ids.append(self.dictionary.word2idx[word])
                idss.append(torch.tensor(ids).type(torch.int64))
            ids = torch.cat(idss)
        return ids

# 设置模型的路径
model_data_filepath = 'data/'
corpus = Corpus(model_data_filepath + 'wikitext-2')
ntokens = len(corpus.dictionary)

# 搭建网络模型
model = LSTMModel(
    ntoken = ntokens,
    ninp = 512,
    nhid = 256,
    nlayers = 5,
)

# 装载预训练的权重
model.load_state_dict(
    torch.load(
        model_data_filepath + 'word_language_model_quantize.pth',
        map_location=torch.device('cpu')
        )
    )
# 将模型切换为推理模式，并打印整个模型
model.eval()
print(model)

# 获取一个随机的输入数值
input_ = torch.randint(ntokens, (1, 1), dtype=torch.long)
hidden = model.init_hidden(1)
temperature = 1.0
num_words = 1000

# 遍历数据集进行前向推理并将结果保存起来
with open(model_data_filepath + 'out.txt', 'w') as outf:
    with torch.no_grad():  # no tracking history
        for i in range(num_words):
            output, hidden = model(input_, hidden)
            word_weights = output.squeeze().div(temperature).exp().cpu()
            word_idx = torch.multinomial(word_weights, 1)[0]
            input_.fill_(word_idx)
            word = corpus.dictionary.idx2word[word_idx]
            outf.write(str(word.encode('utf-8')) + ('\n' if i % 20 == 19 else ' '))
            if i % 100 == 0:
                print('| Generated {}/{} words'.format(i, 1000))

with open(model_data_filepath + 'out.txt', 'r') as outf:
    all_output = outf.read()
    print(all_output)

bptt = 25
criterion = nn.CrossEntropyLoss()
eval_batch_size = 1

# 创建测试数据集
def batchify(data, bsz):
    # 对测试数据集进行分块
    nbatch = data.size(0) // bsz
    # 去掉多余的元素
    data = data.narrow(0, 0, nbatch * bsz)
    # 在bsz批处理中平均划分数据
    return data.view(bsz, -1).t().contiguous()

test_data = batchify(corpus.test, eval_batch_size)

# 获取bath块的输入数据
def get_batch(source, i):
    seq_len = min(bptt, len(source) - 1 - i)
    data = source[i:i+seq_len]
    target = source[i+1:i+1+seq_len].view(-1)
    return data, target

def repackage_hidden(h):
  """
  用新的张量把隐藏的状态包装起来，把它们从历史中分离出来
  """
  if isinstance(h, torch.Tensor):
      return h.detach()
  else:
      return tuple(repackage_hidden(v) for v in h)
# 评估函数
def evaluate(model_, data_source):
    # 打开评估模式
    model_.eval()
    total_loss = 0.
    hidden = model_.init_hidden(eval_batch_size)
    with torch.no_grad():
        for i in range(0, data_source.size(0) - 1, bptt):
        	# 获取测试数据
            data, targets = get_batch(data_source, i)
            # 执行前向推理
            output, hidden = model_(data, hidden)
            hidden = repackage_hidden(hidden)
            output_flat = output.view(-1, ntokens)
            # 获取训练loss
            total_loss += len(data) * criterion(output_flat, targets).item()
    return total_loss / (len(data_source) - 1)

# 初始化动态量化模块
quantized_model = torch.quantization.quantize_dynamic(
    model, {nn.LSTM, nn.Linear}, dtype=torch.qint8
)
print(quantized_model)

def print_size_of_model(model):
    torch.save(model.state_dict(), "temp.p")
    print('Size (MB):', os.path.getsize("temp.p")/1e6)
    os.remove('temp.p')

print_size_of_model(model)
print_size_of_model(quantized_model)

torch.set_num_threads(1)
# 评估模型的运行时间
def time_model_evaluation(model, test_data):
    s = time.time()
    loss = evaluate(model, test_data)
    elapsed = time.time() - s
    print('''loss: {0:.3f}\nelapsed time (seconds): {1:.1f}'''.format(loss, elapsed))

time_model_evaluation(model, test_data)
time_model_evaluation(quantized_model, test_data)