create_prune_retrain_model
功能说明
通道稀疏或4选2结构化稀疏接口，两种稀疏特性每次只能使能一个：
将输入的待稀疏的图结构按照给定的稀疏配置文件进行稀疏处理，在传入的图结构中插入或者替换相关的算子，生成记录稀疏信息的record_file，返回修改后可用于稀疏后训练的torch.nn.module模型。
函数原型
prune_retrain_model = create_prune_retrain_model (model, input_data, config_defination，record_file)
参数说明
参数名 |
输入/返回值 |
含义 |
使用限制 |
|---|---|---|---|
model |
输入 |
待进行稀疏的模型，已加载权重。 |
数据类型：torch.nn.module |
input_data |
输入 |
模型的输入数据。一个torch.tensor会被等价为tuple(torch.tensor)。 |
数据类型：tuple |
config_defination |
输入 |
简易配置文件。 基于retrain_config_pytorch.proto文件生成的简易配置文件prune.cfg， retrain_config_pytorch.proto文件所在路径为：AMCT安装目录/amct_pytorch/proto/retrain_config_pytorch.proto。 retrain_config_pytorch.proto文件参数解释以及生成的prune.cfg简易量化配置文件样例请参见量化感知训练简易配置文件说明。 |
数据类型：string |
record_file |
输入 |
记录稀疏信息的文件路径及名称，记录通道稀疏结点间的级联关系或记录4选2稀疏的节点。 |
数据类型：string |
prune_retrain_model |
返回值 |
修改后可用于稀疏后训练的torch.nn.module模型。 |
默认值：None 数据类型：torch.nn.module |
返回值说明
稀疏后的模型。
函数输出
无。
调用示例
1 2 3 4 5 6 7 8 9 10 11 12 13 14 |
import amct_pytorch as amct # 建立待进行稀疏的网络图结构 model = build_model() model.load_state_dict(torch.load(state_dict_path)) input_data = tuple([torch.randn(input_shape)]) # 调用稀疏模型API record_file = os.path.join(TMP, 'scale_offset_record.txt') cfg_file = './prune_config.cfg' prune_retrain_model = amct.create_prune_retrain_model( model, input_data, cfg_file, record_file) |
