下载
中文
注册

create_compressed_retrain_model

功能说明

静态组合压缩接口,将输入的待静态组合压缩的模型按照给定的组合压缩配置文件进行压缩处理,即将传入的模型先进行稀疏(通道稀疏或者4选2结构化稀疏,二选一),后对模型插入量化相关的算子(数据和权重的量化感知训练层以及searchN的层),生成稀疏和量化因子记录文件record_file(如果配置存在),返回修改后的torch.nn.Module模型。

约束说明

组合压缩配置文件至少存在一个配置:稀疏配置或者量化配置。

函数原型

compressed_retrain_model = create_compressed_retrain_model(model, input_data, config_defination, record_file)

参数说明

参数名

输入/返回值

含义

使用限制

model

输入

PyTorch的model。

数据类型:torch.nn.Module

input_data

输入

模型的输入数据。

数据类型:tuple

config_defination

输入

静态组合压缩简易配置文件。

基于retrain_config_pytorch.proto文件生成的简易配置文件compressed.cfg,

retrain_config_pytorch.proto文件所在路径为:AMCT安装目录/amct_pytorch/proto/retrain_config_pytorch.proto。

retrain_config_pytorch.proto文件参数解释以及生成的compressed.cfg简易配置文件样例请参见量化感知训练简易配置文件说明

数据类型:string

record_file

输入

待记录稀疏和量化因子文件路径及名称。

数据类型:string

compressed_retrain_model

返回值

根据配置文件进行稀疏后(如果配置稀疏),且插入量化相关层(如果配置量化)的torch.nn.Module。

数据类型:torch.nn.Module

返回值说明

静态组合压缩的模型。

函数输出

静态组合压缩记录文件record_file,如果简易配置文件中含有稀疏配置,则在该函数完成后,record_file中含有稀疏记录信息。

调用示例

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
import amct_pytorch as amct
# 建立待进行静态组合压缩的网络
model = build_model()
input_data = tuple([torch.randn(input_shape)])

# 调用静态组合压缩API
record_file = os.path.join(TMP, 'compressed_record.txt')
config_defination = './compressed_cfg.cfg'

compressed_retrain_model = amct.create_compressed_retrain_model(
                                model,
                                input_data,
                                config_defination,
                                record_file)