auto_decomposition

产品支持情况

产品

是否支持

Atlas A2 训练系列产品/Atlas 800I A2 推理产品/A200I A2 Box 异构组件

Atlas 200I/500 A2 推理产品

Atlas 推理系列产品

Atlas 训练系列产品

功能说明

根据用户提供的TensorFlow模型文件,生成张量分解后的模型文件。

函数原型

1
add_ops = auto_decomposition(meta_path, ckpt_path, save_path)

参数说明

参数名

输入/输出

说明

meta_path

输入

含义:待分解的TensorFlow模型定义文件路径,格式为.meta。

数据类型:string

说明:

请确保.meta能被tf.compat.v1.train.import_meta_graph正常加载。

例如:对于使用horovod训练得到的meta文件,必须执行import hovorod.tensorflow才能被auto_decomposition接口调用。

ckpt_path

输入

含义:待分解的TensorFlow模型权重文件路径,该参数为.data-XXXXX-of-XXXXX文件与.index文件的共同前缀路径。

例如上述两个文件分别为path/model-200.data-00000-of-00001和path/model-200.index,

则ckpt_path取值为path/model-200。

数据类型:string

save_path

输入

含义:张量分解后所得文件的保存路径,将生成以该路径为前缀的.data-XXXXX-of-XXXXX等文件。

例如该参数设置为path/model,则生成的.data-XXXXX-of-XXXXX文件与.index等文件为:

path/model.data-XXXXX-of-XXXXX和path/model.index。

数据类型:string

返回值说明

张量分解后新添加卷积层的名字的列表,类型为list。

约束说明

调用示例

1
2
3
4
5
from amct_tensorflow.tensor_decompose import auto_decomposition
auto_decomposition(meta_path='src_path/model.meta',
                   ckpt_path='src_path/model',
                   save_path='decomposed_path/model'
)

落盘文件: