Brevitas项目教程：量化循环神经网络(RNN)与长短期记忆网络(LSTM)的实现

Brevitas项目教程：量化循环神经网络(RNN)与长短期记忆网络(LSTM)的实现

brevitas Brevitas: neural network quantization in PyTorch 项目地址: https://gitcode.com/gh_mirrors/br/brevitas

引言

在深度学习模型部署到资源受限设备的场景中，模型量化已成为必不可少的技术手段。Brevitas作为一个专注于量化神经网络的开源库，在0.8版本中引入了对量化循环神经网络(RNN)和长短期记忆网络(LSTM)的支持。本文将详细介绍如何使用Brevitas实现量化RNN和LSTM，并深入解析其技术细节。

量化RNN实现

基本结构

Brevitas中的QuantRNN类提供了量化RNN的实现，其构造函数如下：

class QuantRNN:
    def __init__(
            self,
            input_size: int,
            hidden_size: int,
            num_layers: int = 1,
            nonlinearity: str = 'tanh',
            bias: Optional[bool] = True,
            batch_first: bool = False,
            bidirectional: bool = False,
            weight_quant=Int8WeightPerTensorFloat,
            bias_quant=Int32Bias,
            io_quant=Int8ActPerTensorFloat,
            gate_acc_quant=Int8ActPerTensorFloat,
            shared_input_hidden_weights=False,
            return_quant_tensor: bool = False,
            dtype: Optional[torch.dtype] = None,
            device: Optional[torch.device] = None,
            **kwargs):

与标准PyTorch RNN相比，QuantRNN增加了四个量化器参数：

• weight_quant: 控制权重张量的量化
• bias_quant: 控制偏置项的量化
• io_quant: 统一控制输入/输出的量化
• gate_acc_quant: 控制门计算输出的量化

量化器共享机制

在多层和双向RNN中，量化器的共享策略有所不同：

1. 权重、门计算和偏置量化器：各层/方向共享相同的量化器定义，但每个层/方向拥有独立的量化器实例

# 不同层/方向的权重量化器实例不同
assert not quant_rnn_0_left_to_right.gate_params.input_weight.weight_quant is quant_rnn_1_right_to_left.gate_params.input_weight.weight_quant

2. 输入输出量化器：所有层和方向共享同一个实例，确保内部拼接的张量使用相同的量化参数

# 所有层/方向共享相同的io_quant实例
assert quant_rnn_0_left_to_right.io_quant is quant_rnn_1_right_to_left.io_quant

权重共享优化

QuantRNN支持shared_input_hidden_weights标志，当设置为True时，双向RNN的两个方向会共享输入到隐藏层的权重，这可以显著减少参数量：

# 单方向RNN参数数量: 600
# 双向RNN参数数量: 1200 
# 共享输入权重的双向RNN参数数量: 1000

量化参数配置

可以通过关键字参数直接修改量化器配置，例如设置4位每通道权重和6位输入输出量化：

quant_rnn_4b = QuantRNN(
    input_size=10, 
    hidden_size=20,
    weight_bit_width=4,
    weight_scaling_per_output_channel=True,
    io_bit_width=6)

量化LSTM实现

基本结构

QuantLSTM的实现与QuantRNN类似，但针对LSTM特有的门控机制进行了优化：

class QuantLSTM:
    def __init__(
            self,
            input_size: int,
            hidden_size: int,
            num_layers: int = 1,
            bias: Optional[bool] = True,
            batch_first: bool = False,
            bidirectional: bool = False,
            weight_quant=Int8WeightPerTensorFloat,
            bias_quant=Int32Bias,
            io_quant=Int8ActPerTensorFloat,
            gate_acc_quant=Int8ActPerTensorFloat,
            shared_input_hidden_weights=False,
            return_quant_tensor: bool = False,
            dtype: Optional[torch.dtype] = None,
            device: Optional[torch.device] = None,
            **kwargs):

门控量化

LSTM包含四个门(输入门、遗忘门、输出门和候选记忆)，每个门的计算都可以独立量化：

# 获取输入门的量化权重
input_gate_weight = quant_lstm_0_left_to_right.gate_params.input_gate.input_weight.quant_weight()

使用注意事项

1. 输入输出格式：

   • 支持形状为(batch, sequence, features)(batch_first=True)或(sequence, batch, features)(batch_first=False)的输入
   • 输出为元组(outputs, hidden_states)

2. 当前限制：

   • 不支持打包的可变长度输入
   • 不支持非批处理输入

3. 量化张量返回： 设置return_quant_tensor=True可以返回包含量化元数据的QuantTensor

outputs, hidden_states = quant_rnn(input_tensor, return_quant_tensor=True)

性能优化建议

1. 量化位宽选择：根据硬件支持选择合适的位宽，通常4-8位在精度和效率间取得良好平衡
2. 共享权重：在双向RNN中考虑使用shared_input_hidden_weights减少参数量
3. 分层量化：对不同层使用不同的量化策略，关键层使用较高位宽

结语

Brevitas提供的QuantRNN和QuantLSTM为在资源受限设备上部署循环神经网络提供了强大支持。通过灵活的量化配置和优化选项，开发者可以在模型大小、计算效率和精度之间找到最佳平衡点。随着量化技术的不断发展，量化循环神经网络将在边缘计算和物联网等领域发挥越来越重要的作用。

brevitas Brevitas: neural network quantization in PyTorch 项目地址: https://gitcode.com/gh_mirrors/br/brevitas