docs/tools/mnncompress.md
:orphan:
mnncompress 是基于主流训练框架TF/Pytorch的模型压缩工具，需要训练数据和对应的训练框架环境） ，提供了包括低秩分解、剪枝、量化等模型压缩算法的实现，并且MNN进一步实现了其中一些需要软件特殊实现的算法（如稀疏计算和量化）的底层计算过程，此工具箱需要配合MNN推理框架来使用。目前提供的能力如下表所示：
| 原始模型格式 | 提供的工具 | 支持的压缩算法类别 |
|---|---|---|
| TensorFlow 1.X | python压缩算法插件（mnncompress） | 低秩分解 |
| 线性过参数化 | ||
| 模型剪枝 | ||
| 离线量化 | ||
| 训练量化 | ||
| 训练权值量化 | ||
| Pytorch | python压缩算法插件（mnncompress） | 低秩分解 |
| 模型剪枝 | ||
| 离线量化 | ||
| 训练量化 | ||
| 训练权值量化 |
各类压缩算法的特点：
| 压缩算法类别 | 特点 | 支持的算法 |
|---|---|---|
| 低秩分解 | 将原始模型进行分解，降低模型计算量、存储量，分解之后的模型仍是一个规整的模型，不需要特殊的软件底层实现；需要进行finetune训练 | Tucker分解，SVD分解 |
| 线性过参数化 | 用于模型从0开始训练的过程中，设计一个小模型，然后使用此算法对模型进行过参数化扩展为一个大模型，提高模型的表达能力，然后基于此大模型进行训练，训练完成之后可以将大模型中的层合并得到原始小模型一样的结构，而且精度和大模型一致 | 改进的 linear over parameterization |
| 模型剪枝 | 将模型的权值矩阵进行稀疏，并进行finetune训练，利用稀疏编码（通道剪枝不需要）压缩模型大小，底层稀疏计算实现加速 | 1x1随机剪枝，1x4block剪枝，通道剪枝 |
| 离线量化 | 将float卷积转换为int8卷积计算，仅需少量校准图片，降低存储量到原始模型的四分之一，降低内存，加速计算（某些模型可能会比float模型慢，因为float的优化方法和int8不同） | EMA，KL，ADMM |
| 训练量化 | 将float卷积转换为int8卷积计算，需要进行训练，可提高量化模型精度，降低存储量到原始模型的四分之一，降低内存，加速计算（某些模型可能会比float模型慢，因为float的优化方法和int8不同） | LSQ，OAQ，WAQ |
| 训练权值量化 | 特点同直接权值量化，但通过mnncompress压缩算法插件实现，因而可以提供更低比特的权值量化，以减少更多的存储量，并提高权值量化之后模型的精度，例如4bit量化情况下，模型大小减小到原来的1/8 | 对称量化 |
| FP16 | 将FP32的权重转换成FP16的类型，可在模型转换时一键完成，模型大小减小为原来的1/2，精度基本无损 | - |
使用pip安装：pip install mnncompress
| 算法类别 | 算法名称 | TensorFlow1.X | Pytorch |
|---|---|---|---|
| 线性过参数化 | linear over parameterization 改进 | ❎ | ✅ |
| 低秩分解 | Tucker分解，SVD分解 | ✅ | ✅ |
| 剪枝 | TaylorFOChannelPruner 结构化通道剪枝 | ✅ | ✅ |
| SIMDOCPruner 1*4半结构化剪枝 | ✅ | ✅ | |
| SNIPLevelPruner 1*1随机非结构化剪枝 | ✅ | ✅ | |
| 训练量化（也可改做离线量化） | LSQ（学习量化参数） | ✅ | ✅ |
| OAQ（overflow aware quantization） | ✅ | ❎ | |
| WAQ（winograd aware quantization） | ❎ | ✅ | |
| 权值量化（也可不训练，查看直接权值量化精度） | 对称量化，非对称量化 | ✅ | ✅ |
通过mnncompress中的模型压缩插件对原始模型进行压缩，然后得到一个原始的float模型，以及该压缩算法所需要的额外的压缩参数（例如对于量化算法来说，需要各个卷积层的scale和zeropoint等信息）。最后将这个float模型和模型压缩参数文件输入到MNN转换器中，得到最终的MNN压缩模型。
在mnncompress的文档中通常命名为“compress_params.bin”，是一个protobuf序列化之后的二进制文件。每执行一个压缩算法都需要保存一个对应的模型压缩参数文件，如果有算法上的叠加（例如剪枝+量化的叠加），那么需要将算法对应的API接口中的append参数设置为True（例如剪枝+量化的叠加，则需要在量化算法保存压缩参数文件时，将对应的接口append参数设置为True，将量化算法所需要的参数append到剪枝算法的模型压缩文件上去），这样MNN模型转换器才会知道你的模型经过了哪些优化，才能进行正确的转换。
以下模型均来自torchvision预训练模型，采用mnncompress中LSQ的offline模式进行量化，batch size为64，用10个batch，640张训练图片进行量化，表中标*的模型表示量化时跳过了第一层或者前两层不量化，以提升精度。测速均采用华为P20 Pro 单线程。
| 模型 | 原始模型指标 | 压缩模型指标 | ARMV7 (ms) | ARMV8 (ms) |
|---|---|---|---|---|
| ResNet-18 | 69.758%，45M | 69.740%，12M | 196.6 --> 208.8 | 187.6 --> 167.0 |
| ResNet-50 | 76.130%，98M | 76.030%，25M | 606.0 --> 470.5 | 550.8 --> 379.6 |
| SqueezeNet 1.0 | 58.092%，4.8M | 57.800%，1.3M | 122.1 --> 104.1 | 120.8 --> 88.3 |
| ShuffleNet V2 x1.0 | 69.362%，8.7M | 68.616%，2.3M | 33.3 --> 33.8 | 29.0 --> 26.0 |
| MobileNet V2 | 71.878%，14M | 71.150%，3.5M | 69.1 --> 50.2 | 62.2 --> 42.1 |
| *MobileNet V3 Large | 74.042%，21M | 73.030%，5.4M | 68.6 --> 63.1 | 64.9 --> 50.1 |
| MNASNet 1.0 | 73.456%，17M | 72.692%，4.3M | 70.6 --> 51.2 | 63.0 --> 42.9 |
| *EfficientNet-B0 | 77.692%，21M | 70.486%，5.3M | 134.3 --> 113.0 | 128.7 --> 100.4 |
| EfficientNet-B1 | 78.642%，30M | 73.546%，7.8M | 199.5 --> 166.9 | 185.2 --> 145.9 |
| regnet_x_400mf | 72.834%，21.3M | 72.660%，8.0M | 83.2 --> 67.6 | 75.1 --> 58.0 |
训练量化用来提升量化的精度，其速度和离线量化版本模型一致。建议优先使用离线量化方法，精度不够的情况下再使用训练量化方法。以下数据中的标准模型来自torchvision预训练模型，有些模型如efficientnet的训练成本较大，故未给出训练量化版本。用户自己训练的模型知道训练参数，结果一般可以更好。
| 模型 | 原始模型指标 | 压缩模型指标 | 备注 |
|---|---|---|---|
| ResNet-18 | 69.758%，45M | 69.840%，12M | 训练一个epoch即可 |
| MobileNet V2 | 71.878%，14M | 71.762%，3.5M | github: pytorch/vision reference/classification，复现命令：torchrun --nproc_per_node=8 train.py --model mobilenet_v2 --data-path /mnt/data/ --epochs 100 --lr 0.01 --wd 0.00004 --lr-step-size 30 --lr-gamma 0.1 --pretrained --quant --sync-bn -b 128 |
| *MobileNet V3 Large | 74.042%，21M | 73.924%，5.4M | 跳过第一层，github: pytorch/vision reference/classification，复现命令：torchrun --nproc_per_node=8 train.py --opt rmsprop --auto-augment imagenet --random-erase 0.2 --model mobilenet_v3_large --data-path /mnt/data/ --epochs 100 --batch-size 128 --lr 0.01 --wd 0.00001 --lr-step-size 30 --lr-gamma 0.1 --pretrained --quant --sync-bn |
| 模型 | 原始模型指标 | 压缩模型指标 | ARMV7 (ms) | ARMV8 (ms) |
|---|---|---|---|---|
| MobileNet V2 | 71.878%，14M | 71.272%，2.8M，50% SIMDOC稀疏 | 69.1 --> 64.9 | 62.2 --> 58.8 |
| *MobileNet V3 Large | 74.042%，21M | 73.568%，4.2M 50% SIMDOC稀疏 | 68.6 --> 66.5 | 64.9 --> 62.3 |
| 模型 | 原始模型指标 | 压缩模型指标 | ARMV7 (ms) | ARMV8 (ms) |
|---|---|---|---|---|
| MobileNet V2 | 71.878%，14M | 69.874%，11M | 69.1 --> 60.6 | 62.2 --> 54.6 |
| MobileNet V3 Large | 74.042%，21M | 72.748%，18M | 68.6 --> 59.7 | 64.9 --> 55.7 |
| 模型 | 压缩方案 | 原始模型指标 | 压缩模型指标 |
|---|---|---|---|
| 神经渲染relight模型 | 低秩分解 | 30.1，138M，238ms | 29.98，58M，169ms / 29.68，17M，116ms |
| 语音识别AOA V3 | EMA训练量化 | cer 18.3，50.6M | cer 18.26，18.7M |
from mnncompress.pytorch import LOP
# 定义你的模型结构，并初始化
model = Net()
# 对模型进行线性过参数化
lop = LOP(model)
# 扩大8倍，指定模型压缩参数文件，更多参数查看api文档
expand_model = lop.linear_expand_layers(8, "compress_params.bin")
# 使用线性过参数化之后的模型进行训练
train_and_evaluate(expand_model, data)
# 保存模型之前，将过参数化的模型合并，然后保存合并之后的模型 merged_model
merged_model = lop.linear_merge_layers()
LOP(model) # model为原始模型nn.Module实例
linear_expand_layers：
linear_expand_layers(expand_rate, compress_params_file, add_batchnorm=True, add_bypass=True, append=False)
'''
将原始模型进行线性过参数化
参数：
expand_rate: int，线性过参数化的倍数，一般取2，4，8，16等
compress_params_file: str，模型压缩参数文件名
add_batchnorm: bool，是否在线性过参数化时使用BN
add_bypass: bool，是否在线性过参数化时添加bypass（残差连接）
append: bool，是否将线性过参数化算法的参数追加到compress_params_file中去，为False，则将覆盖compress_params_file
返回值：
线性过参数化之后的模型，深拷贝，和原始model不共享内存
'''
linear_merge_layers
linear_merge_layers()
'''
将过参数化之后的模型合并为原始模型结构
参数：
无参数
返回值：
合并之后的模型，深拷贝，和过参数化模型不共享内存
'''
from mnncompress.pytorch import low_rank_decompose
# 加载已经训练好的float模型
model = Net()
model.load_state_dict(torch.load("ori_float_model.pt"))
# 将原始模型分解，然后用分解之后的模型进行finetune训练即可
model = low_rank_decompose(model, "compress_params.bin")
low_rank_decompose(model, compress_params_file, skip_layers=[""], align_channels=8,
in_place=False, tucker_minimal_ratio=0.25,
reserved_singular_value_ratio=0.5, append=False)
'''
参数：
model: nn.Module实例，训练好的float模型
compress_params_file: str, MNN模型压缩参数文件名
skip_layers: List[str], 跳过不分解层的名字，需要是nn.Conv2d或者nn.Linear类型，如["features.conv1",]
align_channels: int, 分解之后进行通道对齐的倍数
in_place: 分解时是否使用原模型内存空间，若为False，则分解之前会对原模型进行深拷贝
tucker_minimal_ratio: float 0~1, 卷积层tucker分解保留的最低通道数比例
reserved_singular_value_ratio: svd分解保留的特征值之和占总特征值之和的比例
append: bool, 是否将低秩分解算法的参数追加到compress_params_file中去，为False，则将覆盖compress_params_file
返回值：
分解之后的模型，nn.Module实例
'''
from mnncompress.pytorch.SNIP_level_pruner import SNIPLevelPruner
from mnncompress.pytorch.SIMD_OC_pruner import SIMDOCPruner
Pruner = SIMDOCPruner
# 你的模型代码
class Net(nn.Module):
pass
model = Net()
# 加载已经训练好的模型
model.load_state_dict(torch.load("ori_model.pt"))
model.to(device)
# 将模型进行转换，并使用转换后的模型进行训练，测试
# 更多配置请看API部分
pruner = SIMDOCPruner(model, total_pruning_iterations=1, sparsity=0.6, debug_info=False)
def train(model, data, optimizer, pruner)
model.train()
for d, t in data:
optimizer.zero_grad()
output = model(d)
loss = F.nll_loss(output, t)
loss.backward()
optimizer.step()
# step之后调用pruner的剪枝方法
pruner.do_pruning()
# 获取当前剪枝比例
print(pruner.current_prune_ratios())
for epoch in range(1, epochs + 1):
train(model, data, optimizer, pruner)
test(model, data)
# 保存模型
model.eval()
torch.save(model.state_dict(), "pruned_model.pt")
x = torch.randn(input_shape).to(device)
torch.onnx.export(model, x, "pruned_model.onnx")
# 保存MNN模型压缩参数文件，应在剪枝完毕之后进行，建议在保存模型时调用
pruner.save_compress_params("compress_params.bin", append=False)
mnnconvert --modelFile pruned_model.onnx --MNNModel pruned_model.mnn --framework ONNX --bizCode MNNTest --compressionParamsFile compress_params.bin --weightQuantBits 8
Pruner(model, sparsity=0.5, total_pruning_iterations=1, config_file=None, debug_info= False,
prune_finetune_iterations=0, max_prune_ratio=0.99,
align_channels=4 # only for 'TaylorFOChannelPruner'
):
'''
参数:
model：Module，pytorch模型
sparsity：float，0~1，模型总体剪枝比例
total_pruning_iterations：int，指定多少个iteration将指定比例的连接剪掉
config_file：str，可先运行一次Pruner的do_pruning方法，得到搜索出来的剪枝比例yml文件之后，
然后微调此yml文件中的各层剪枝比例，将修改后的文件名通过此参数传入，可控制各层剪枝比例
debug_info：是否输出debug信息
prune_finetune_iterations: int, 指定剪枝进行一步，finetune多少步，如指定为99，则剪枝一次，
finetune99次，在finetune的99次过程中剪枝mask不变
max_prune_ratio: float, 指定各层最大剪枝比例
align_channels: int, 仅通道剪枝算法TaylorFOChannelPruner有此参数，用于指定剪枝之后对齐的通道数
方法和属性:
do_pruning(result_file = "found_prune_ratios.yml")：进行一次剪枝，在训练optimizer.step之后调用，搜索到的各层剪枝比例将保存到此文件中
current_prune_ratios()：返回当前各层剪枝比例
save_compress_params(filename, append=False): 保存MNN模型压缩参数到filename文件
append 表示是否将剪枝参数信息append到传入的文件中，如果为false，则将新建或覆盖该文件，
append=True 表示剪枝之前还做了其他模型压缩操作，filename文件中已存在相关压缩信息
'''
from mnncompress.pytorch import WeightQuantizer
# 你的模型代码
class Net(nn.Module):
pass
model = Net()
# 加载已经训练好的模型，可以是剪枝之后的
model.load_state_dict(torch.load("pruned_model.pt"))
# 将模型进行转换，并使用转换后的模型进行训练，测试
# 更多配置请看API部分
quantizer = WeightQuantizer(model, bits=8)
quant_model = quantizer.convert()
quant_model.to(device)
for epoch in range(1, epochs + 1):
# 每次训练之前加上这一句，准备好量化训练图
quantizer.resume_wq_graph()
train(quant_model, data, optimizer)
test(quant_model, data)
if 触发模型保存条件：
# 保存模型之前去掉插入的节点，恢复原模型结构
quantizer.strip_wq_ops()
# 保存模型，注意index，即模型和保存MNN模型压缩参数文件是一一对应的
quant_model.eval()
torch.save(quant_model.state_dict(), "quant_model_index.pt")
x = torch.randn(input_shape).to(device)
torch.onnx.export(quant_model, x, "quant_model_index.onnx")
# 保存MNN模型压缩参数文件，如果进行量化的模型有剪枝，
# 请将剪枝时生成的MNN模型压缩参数文件 "compress_params.bin" 文件在下方传入，并将 append 设置为True
quantizer.save_compress_params("compress_params_index.bin", append=False)
--weightQuantBits指定量化比特数：mnnconvert --modelFile quant_model_index.onnx --MNNModel weight_quant_model.mnn --framework ONNX --bizCode MNNTest --compressionParamsFile compress_params_index.bin --weightQuantBits 8
WeightQuantizer(model, bits = 8, debug_info = False, mode = 'symmetric')
'''
参数:
model：Module，pytorch模型
bits：int，指定权值量化比特数
debug_info：bool，是否输出debug信息
mode："symmetric"或"asymmetric"，采用对称或者非对称量化，目前仅支持对称量化
方法和属性:
convert()：返回用于训练量化的模型
strip_wq_ops()：去掉插入的权值量化op，恢复原始模型结构
resume_wq_graph()：strip_wq_ops保存模型之后，如果还要继续量化训练需加上这一句，以恢复量化训练图，由此支持边训练边保存
save_compress_params(filename, append=False)：
用于保存MNN转换时需要用的模型压缩信息
filename：str，MNN模型压缩参数将保存到这个文件名指定的文件中
append：bool，是否将量化参数追加到filename文件中。如果进行量化的模型有剪枝，请将剪枝时通过save_compress_params生成的剪枝信息文件通过此参数传入，并将 append 设置为True
'''
from mnncompress.pytorch import LSQQuantizer
Quantizer = LSQQuantizer
# 你的模型代码
class Net(nn.Module):
pass
model = Net()
# 加载已经训练好的模型，可以是剪枝之后的
model.load_state_dict(torch.load("pruned_model.pt"))
# 将模型进行转换，并使用转换后的模型进行训练，测试
# retain_sparsity=True表示待量化的float模型是稀疏模型，希望叠加量化训练
# 更多配置请看API部分
# 注意在model还在cpu上，任何分布式还未生效前调用
quantizer = Quantizer(model, retain_sparsity=False)
quant_model = quantizer.convert()
# 单机或分布式，根据你已有的代码来
quant_model.to(device)
for epoch in range(1, epochs + 1):
# 每次训练之前加上这一句，准备好量化训练图
quantizer.resume_qat_graph()
train(quant_model, data, optimizer)
test(quant_model, data)
if 触发模型保存条件：
# 保存模型之前去掉插入的节点，恢复原模型结构
quantizer.strip_qat_ops()
# 保存模型，注意index，即模型和保存MNN模型压缩参数文件是一一对应的
quant_model.eval()
torch.save(quant_model.state_dict(), "quant_model_index.pt")
x = torch.randn(input_shape).to(device)
torch.onnx.export(quant_model, x, "quant_model_index.onnx")
# 保存MNN模型压缩参数文件，如果进行量化的模型有剪枝，
# 请将剪枝时生成的MNN模型压缩参数文件 "compress_params.bin" 文件在下方传入，并将 append 设置为True
quantizer.save_compress_params("quant_model_index.onnx", "compress_params_index.bin", append=False)
mnnconvert --modelFile quant_model.onnx --MNNModel quant_model.mnn --framework ONNX --bizCode MNNTest --compressionParamsFile compress_params.bin
LSQQuantizer的mode设置为offline（查看下方API），其他步骤和上述一致，此时该工具即成为离线量化工具，直接灌入训练数据即可完成离线量化：
# 注释掉训练部分中的这一句即可
# optimizer.step()
LSQQuantizer(model, skip_quant_layers = [], bits = 8, debug_info = False, mode = 'online', retain_sparsity=False)
'''
参数:
model：Module，pytorch模型
skip_quant_layers：[str,]，不进行量化的module的名字，可通过dict(model.named_modules()).keys()获得，如['model.conv1', 'model.conv2']，需要是nn.Conv2d或者nn.Linear
bits：int，指定量化比特数
debug_info：bool，是否输出debug信息
mode："online"或"offline"，训练量化时选online，离线量化时选offline
retain_sparsity: bool，为True表示待量化的float模型是稀疏模型，希望叠加量化训练
方法和属性:
convert()：返回用于训练量化的模型
strip_qat_ops()：去掉插入的训练量化op，恢复原始模型结构
resume_qat_graph()：strip_qat_ops保存模型之后，如果还要继续量化训练需加上这一句，以恢复量化训练图，由此支持边训练边保存
save_compress_params(onnx_inference_model_file, filename, append=False)：
用于保存MNN转换时需要用的模型压缩信息
onnx_inference_model_file：str，去掉插入的训练量化节点之后保存的onnx推理模型名字
filename：str，MNN模型压缩参数将保存到这个文件名指定的文件中
append：bool，是否将量化参数追加到filename文件中。如果进行量化的模型有剪枝，请将剪枝时通过save_compress_params生成的剪枝信息文件通过此参数传入，并将 append 设置为True
'''
from mnncompress.tensorflow import get_op_weight_values, low_rank_decompose
# first, 构建前向网络模型
build_model()
# 新建一个临时session，此session会在分解之后关闭掉
temp_sess = tf.Session()
# 给原模型中的变量用checkpoint赋值
saver = tf.train.Saver()
saver.restore(sess, 'save/model.ckpt')
# 获得图中的权值numpy数值
get_op_weight_values(temp_sess, "model_weights.npy")
temp_sess.close()
# 将模型图进行分解
low_rank_decompose(tf.get_default_graph(), "model_weights.npy", "compress_params.bin")
# 此时图已经分解，构建optimizer，正常训练即可
optimizer = ...
# 恢复模型中未分解层的参数
saver = tf.train.Saver()
saver.restore(sess, 'save/model.ckpt')
# 初始化分解之后的层的参数
initializers = tf.global_variables_initializer()
get_op_weight_values(sess, npy_file_name)
'''
参数：
sess: 临时tf.Session，获取到权值numpy数值之后会被关掉
npy_file_name: str, 模型中权值的数值会被保存到此npy文件中，用于low_rank_decompose分解
返回值：
无
'''
low_rank_decompose(graph, weight_npy_file, compress_params_file, skip_layers=[""], align_channels=8,
tucker_minimal_ratio=0.25,
reserved_singular_value_ratio=0.5, append=False)
'''
参数：
graph: 模型前向图
weight_npy_file: str, get_op_weight_values函数得到的模型权值npy文件名
compress_params_file: str, MNN模型压缩参数文件名
skip_layers: List[str], 跳过不分解层的名字，需要是nn.Conv2d或者nn.Linear类型，如["features.conv1",]
align_channels: int, 分解之后进行通道对齐的通道数
tucker_minimal_ratio: float 0~1, 卷积层tucker分解保留的最低通道数比例
reserved_singular_value_ratio: svd分解保留的特征值之和占总特征值之和的比例
append: bool, 是否将低秩分解算法的参数追加到compress_params_file中去，为False，则将覆盖compress_params_file
返回值：
无返回值，此时当前graph已经被分解
'''
import tensorflow as tf
from mnncompress.tensorflow import SNIPLevelPruner, SIMDOCPruner, TaylorFOChannelPruner
from mnncompress.tensorflow import SensitivityAnalyzer
# 构图
...
# 恢复checkpoint
sess = tf.Session()
sess.run(tf.global_variables_initializer())
saver = tf.train.Saver()
saver.restore(sess, 'original_model/model.ckpt')
# 模型中用到的优化器，例如，optimizer = tf.train.AdamOptimizer(1e-4)
optimizer = ...
# 获取（gradient，variable）pair
gradients_vars = optimizer.compute_gradients(loss)
# 模型更新op
train_step = optimizer.apply_gradients(gradients_vars, global_step=tf.train.get_or_create_global_step())
from mnncompress.tensorflow import SensitivityAnalyzer
analyzer = SensitivityAnalyzer(gradients_vars)
for i in range(10):
batch_x, batch_y = mnist.train.next_batch(50)
batch_x = np.reshape(batch_x,(-1, 28, 28,1))
feed_dict={x_PH: batch_x, labels: batch_y}
# 计算各参数的梯度
gv_numpy = sess.run(gradients_vars, feed_dict=feed_dict)
analyzer.analyze(gv_numpy)
analyzer.save_sensitivity_npy("model_sens.npy")
target_sparsity = 0.6
total_prune_iterations = 1
# 可选SNIPLevelPruner, SIMDOCPruner, TaylorFOChannelPruner等算法，更多配置参数参见下方api文档
pruner = SNIPLevelPruner("model_sens.npy", sess.graph, target_sparsity, total_prune_iterations, debug_info=True)
for data in training_dataset:
feed_dict = {x: batch_x, labels: batch_y}
sess.run([train_step], feed_dict=feed_dict)
# 反向更新之后，执行剪枝，测试时不需要
pruner.do_pruning(sess)
pruner.save_compress_params("compress_params.bin", sess)
output_graph_def = tf.graph_util.convert_variables_to_constants(sess, sess.graph_def, output_node_names=output_names)
f = open("frozen.pb", 'wb')
f.write(output_graph_def.SerializeToString())
mnnconvert --modelFile frozen.pb --MNNModel quant_model.mnn --framework TF --bizCode AliNNTest --compressionParams compress_params.bin --weightQuantBits 8
Pruner(sensitivity_npy_file, graph=None, sparsity=0.5, total_pruning_iterations=1, config_file=None, debug_info= False,
prune_finetune_iterations=0, max_prune_ratio=0.99,
prune_ratio_result_file="算法名_found_prune_ratios.yml",
align_channels=4 # TaylorFOChannelPruner特有参数
)
'''
参数:
sensitivity_npy_file: 敏感度分析得到的npy文件
graph: tensorflow计算图，如为None，则默认使用tf.get_default_graph()来获取计算图
sparsity: float，0~1， 模型的总体剪枝比例
total_pruning_iterations: int， 多少个iteration内将指定比例的参数稀疏掉
config_file: str，可先运行一次Pruner的do_pruning方法，得到搜索出来的剪枝比例yml文件，
然后微调此yml文件中的各层剪枝比例，将修改后的文件名通过此参数传入，可控制各层剪枝比例
debug_info: bool，是否输出debug信息
prune_finetune_iterations: int，剪枝一次，finetune此参数指定步数，然后再剪枝一次。
如果total_pruning_iterations设为10，prune_finetune_iterations设为99，则1000个iteration之后才执行完毕剪枝过程
max_prune_ratio: 0~1，控制各层最大的剪枝比例
prune_ratio_result_file: str，指定搜索得到的剪枝比例保存的文件名
align_channels: int，指定通道剪枝之后对齐的通道数
方法和属性:
do_pruning(sess): 训练时每更新一次模型参数之后，执行剪枝步骤，测试时不需要
save_compress_params(filename, sess, append=False): 保存剪枝参数信息到filename文件，sess为维护训练量化图的session,
append 表示是否将剪枝参数信息append到传入的proto文件中，如果为false，则将新建或覆盖该文件，
append=True 表示剪枝之前还做了其他模型压缩操作，filename文件中已存在相关压缩信息
'''
from mnncompress.tensorflow.weight_quantizer import WeightQuantizer, strip_wq_ops
# 构建 前向 网络模型
build_model_architecture()
# 在定义反向计算之前构建weight quantizer，向图中插入相关节点
graph = tf.get_default_graph()
weight_quantizer = WeightQuantizer(graph, bits=4, debug_info=True)
# 定义optimizer
opt = ...
# 恢复原模型中的变量
saver = tf.train.Saver()
sess = tf.Session()
# 给原模型中的变量用checkpoint赋值
saver.restore(sess, 'save/model.ckpt')
# 训练之前初始化相关变量，放在restore之后
weight_quantizer.init(sess)
for data in training_dataset:
feed_dict = {x: batch_x, labels: batch_y}
sess.run([train_step], feed_dict=feed_dict)
weight_quantizer.update(sess)
# 去掉插入的权值量化算子
strip_wq_ops()
output_graph_def = tf.graph_util.convert_variables_to_constants(sess, sess.graph_def, output_node_names=output_names)
f = open("frozen.pb", 'wb')
f.write(output_graph_def.SerializeToString())
mnnconvert --modelFile frozen.pb --MNNModel weight_quant_model.mnn --framework TF --bizCode MNNTest --compressionParams compress_params.bin --weightQuantBits 8
WeightQuantizer(graph, bits=8, debug_info=False)
'''
参数:
graph: tensorflow模型图
bits: 权值量化的bit数
debug_info: bool，是否输出debug信息
方法和属性:
init(sess): 训练之前初始化内部用到的相关变量
update(sess): 更新内部状态
save_compress_params(filename, append=False)：
用于保存MNN转换时需要用的模型压缩信息
filename：str，MNN模型压缩参数将保存到这个文件名指定的文件中
append：bool，是否将量化参数追加到filename文件中。如果进行量化的模型有剪枝，请将剪枝时通过save_compress_params生成的剪枝信息文件通过此参数传入，并将 append 设置为True
'''
'''
去掉权值量化相关算子
'''
from mnncompress.tensorflow import EMAQuantizer, strip_MNN_QAT_ops, get_qat_state_placeholder, MNN_QAT_variables
from mnncompress.tensorflow import LSQQuantizer, strip_MNN_QAT_ops, get_qat_state_placeholder, MNN_QAT_variables
from mnncompress.tensorflow import OAQQuantizer, strip_MNN_QAT_ops, get_qat_state_placeholder, MNN_QAT_variables
QATQuantizer = EMAQuantizer # 也可以改为使用LSQQuantizer或者OAQQuantizer
# first, 构建前向网络模型
build_model()
# second, 插入fake quant节点，注意此句位置应在构建optimizer之前
# 更多配置详见API部分
quantizer = QATQuantizer(tf.get_default_graph())
# 构建optimizer，反向网络
opt = ...
# 只恢复原模型中的变量
saver = tf.train.Saver(var_list=[v for v in tf.all_variables() if v not in MNN_QAT_variables()])
sess = tf.Session()
# 初始化所有变量，包括MNN训练量化中用到的变量
sess.run(tf.global_variables_initializer())
# 给原模型中的变量用checkpoint赋值
saver.restore(sess, 'save/model.ckpt')
# 注意：该saver只维护了原来的模型的变量，要保存插入fake quant的模型，需要新建一个saver
new_saver = tf.train.Saver()
quantizer.init(sess)
quant_state = quantizer.is_training # 或者 quant_state = get_qat_state_placeholder()
for data in dataset:
feed_dict = {x: batch_x, labels: batch_y, quant_state: True}
sess.run(train_step, feed_dict=feed_dict)
quantizer.update(sess)
# 如果先进行了剪枝训练，那么需要将剪枝中得到的 "compress_params.bin" 文件在下方传入，并设置 append=True
quantizer.save_compress_params("compress_params.bin", sess, append=False)
from mnncompress.tensorflow import EMAQuantizer, strip_MNN_QAT_ops
# 保存frozen pb之前去掉（绕过）插入的量化节点，传入 strip_MNN_QAT_ops 的graph参数
# 也可以是从checkpoint读取出来的graph
strip_MNN_QAT_ops(sess.graph)
output_graph_def = tf.graph_util.convert_variables_to_constants(sess, sess.graph_def, output_node_names=output_names)
f = open("frozen.pb", 'wb')
f.write(output_graph_def.SerializeToString())
mnnconvert --modelFile frozen.pb --MNNModel quant_model.mnn --framework TF --bizCode AliNNTest --compressionParamsFile compress_params.bin
quantizer.save_compress_params("compress_params.bin", sess)
strip_MNN_QAT_ops(sess.graph)
output_graph_def = tf.graph_util.convert_variables_to_constants(sess, sess.graph_def, output_node_names=output_names)
f = open("frozen.pb", 'wb')
f.write(output_graph_def.SerializeToString())
Quantizer(graph=None, is_training = None, skip_quant_layers = [], skip_quant_flag = [], bits = 8,
min_clamp_value = 0, # 此参数仅OAQQuantizer有
debug_info = False, retain_sparsity=False)
'''
参数:
graph: 需要进行训练量化的模型图
is_training: None or Python bool，为None时，Quantizer的is_training属性是一个placeholder，用于控制量化状态，需要在feed_dict中额外传值；为True或False时，则和用户代码中的训练或测试flag一样，is_training是一个constant，不需要在feed_dict中额外赋值
skip_quant_layers: [str], 指定图中跳过量化的op的名字
skip_quant_flag: [str]，跳过op名字中含有指定字符串的op的量化，例如想要跳过模型中整个decoder1的量化，这些op的名字中都含有"decoder1"，那么在此参数中指定["decoder1"]即可
bits: int, 指定量化比特数，目前仅支持8bit量化
min_clamp_value: int, OAQQuantizer特有参数，用于指定最小的量化范围，例如最低量化为6bit，则此数值设置为2^(6-1) - 1 = 31
debug_info: bool，是否输出debug信息
retain_sparsity: bool, 是否保持量化原始float模型的稀疏结构不变，剪枝之后叠加量化时，需要将此参数设为True
'''
# 为获得相关层的名字，我们提供了一个帮助函数：
import mnncompress.tensorflow.helpers as helper
helper.get_op_name_of_type(graph, ['Conv2D', 'DepthwiseConv2dNative', 'MatMul'])
'''
方法和属性:
is_training: property，返回Quantizer内部的 MNN_QAT_is_training 这个 placeholder
strip_qat_ops(): 保存froze pb前去掉（绕过）插入的训练量化节点
save_compress_params(filename, sess, append=False): 保存量化参数信息到filename文件，sess为维护训练量化图的session,
append 表示是否将量化参数信息append到传入的proto文件中，如果为false，则将新建或覆盖该文件，
append=True 表示量化之前还做了其他模型压缩操作，如剪枝，filename文件中已存在相关压缩信息
init(sess): 初始化Quantizer内部信息
update(sess): 更新内部需要的信息
'''
strip_MNN_QAT_ops(graph)
'''
此函数用于保存froze pb前去掉（绕过）插入的训练量化节点
如果训练和测试分开来做，意味着测试时是读取训练得到的checkpoint然后进行测试的
此时图中已经有了训练量化的fake quant节点，但无法使用Quantizer维护的原图信息进行
保存froze pb前去掉（绕过）插入的训练量化节点 的功能
此时可以使用该函数来达到去掉（绕过）fake quant节点的目的
'''
get_qat_state_placeholder(graph=None)
# 用于在无quantizer对象的情况下获取MNN训练量化工具内部标志训练或者测试的placeholder
MNN_QAT_variables()
# 用于获取MNN训练量化特有的量化variable集合
TensorFlow推荐使用tf.train.MonitoredSession，tf.train.MonitoredTrainingSession进行训练，并使用tf.train.SessionRunHook进行训练流程控制。为方便此类用户使用，我们提供了相关压缩工具的hook，初始化hook之后，即可使用。 这些hook的构造方法和使用的压缩算法文档中一致，在此不再详述。 注意，有些hook需要在构造反向图之前创建，例如量化的hook。目前提供的hook有：