Distiller
是一个开源的Python
软件包,用于神经网络压缩研究。网络压缩可以减少神经网络的内存占用,提高推理速度并节省能源。Distiller
提供了一个PyTorch
环境,用于对压缩算法进行原型设计和分析。
主要功能:
github地址:https://github.com/NervanaSystems/distiller 支持PyTorch1.x(在Pytorch 1.1+PyThon3.6上测试成功)。
使用手册:https://nervanasystems.github.io/distiller/ 具体安装见github。
distiller量化模块的使用以及实例代码位于distiller/quantization和examples/classifier_compression下。
量化模块主要包含的功能如下:
下面主要对量化模型的使用进行简单说明,分别对工具自带模型和自己的模型进行量化。
distiller自带一些测试实例如ResNet20+cifar-10, ResNet50+ImageNet等
找到compress_classifier.py文件,如下:
在此文件目录下运行命令行环境,输入以下指令:
$ python3 compress_classifier.py -a resnet20_cifar ../../../data.cifar10 --resume ../ssl/checkpoints/checkpoint_trained_dense.pth.tar --quantize-eval --evaluate
参数解释:-a 表示模型名称(这里是工具自带的模型名称,其他的如resnet32_cifar, resnet44_cifar, resnet56_cifar等等 cifar的模型代码文件位于distiller/models/cifar10/resnet_cifar.py)
../../../data.cifar10
是数据集路径
--resume 表示保存的模型的路径(一般是FP32即将被量化的模型),此参数现在已经更改为--resum-from
--quantize-eval 表示是否对模型量化并进行保存(只有当--evaluate被设置时才有用)
--evaluate 表示是否在测试集上进行测试
具体的其他参数参看distiller/apputils/image_classifier.py文件和distiller/quantization/range_linear.py文件以及github上参数解释。
运行时会对模型进行压缩,然后在测试集上测试,打印出top1和top5以及loss,运行结束后量化模型会保存在logs下。
import logging
from distiller.data_loggers import *
import DataLoader
import distiller.apputils.image_classifier as ic
import torch
from distiller.quantization.range_linear import PostTrainLinearQuantizer
from opt import args
import distiller.apputils as apputils
from distiller.apputils.image_classifier import test
import torch.nn as nn
import os
from your_model import your_model
model_path = ‘你自己的模型’
model = your_model()
model.load_state_dict(model_path)
model = model.cuda()
model.eval()
quantizer = PostTrainLinearQuantizer(model)
quantizer.prepare_model(torch.rand([你的模型输入的大小如[1, 3, 32, 32]]))
apputils.save_checkpoint(0, 'my_model', model, optimizer=None,
name='quantized',
dir='quantization_model')