昇腾社区首页
中文
注册

Stage

产品支持情况

产品

是否支持

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

Atlas 200I/500 A2 推理产品

Atlas 推理系列产品

Atlas 训练系列产品

Atlas 200/300/500 推理产品

功能说明

设置自定义Pass执行阶段。

函数原型

1
PassRegistrationData &Stage(const CustomPassStage stage)

参数说明

参数名

输入/输出

说明

stage

输入

表示执行阶段。

  • kBeforeInferShape:(默认值)框架入口处InferShape之前执行。
  • kAfterInferShape:InferShape之后执行。

    如果自定义Pass在InferShape之后执行,Pass中需要保证改图之后shape的连续性,可以通过InferShapeAndType接口保证:

    1
    2
    3
    4
    5
    6
    7
        // 1. 获取输入节点node1的输出描述
    	TensorDesc output_desc;
        node1.GetOutputDesc(0, output_desc);
        // 2. 更新当前节点node2的输入描述
        node2.UpdateInputDesc(0, output_desc);
        // 3. 当前节点node2推导InferShape
        operator2.InferShapeAndType();
    

返回值说明

返回PassRegistrationData对象。

约束说明