LookupSubgraphSlicerHook

该Hook用于在查表KeyTensor的子图中查找指定类型算子,然后将查找到的指定类型算子及其最小依赖子图切换到CPU预取阶段执行。如果没有找到目标类型的算子,不会执行切分操作。

该Hook的使用场景是NPUEstimator模式,启用自动改图功能。该Hook的调用时机需要在自动改图的GraphModifierHook之前。

参数名

参数类型

参数说明

必选/可选

op_types

list[str]

指定需要切分的算子类型,目前仅支持列表格式。

必选

使用示例

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
from mx_rec.graph import LookupSubgraphSlicerHook, GraphModifierHook


def input_fn():
    """
    用户自定义Estimator输入函数。
    """

lookup_slicer_hook = LookupSubgraphSlicerHook(op_types=["StringToNumber"] )
modifier_hook = GraphModifierHook(modify_graph=params.modify_graph)
hooks_list = [lookup_slicer_hook, modifier_hook]

est = NPUEstimator(...)
est.train(input_fn=lambda: input_fn, hooks=npu_hooks_append(hooks_list))