Class BlockBuilder

ppq/quantization/algorithm/training.py:191–315

Network Block Builder will cut your graph into small blocks(subgraphs) Each block will have exact 1 input operation and 1 output operation. Besides, there is a parameter 'limit' to control the size of each block.

class BlockBuilder:
    """
    Network Block Builder will cut your graph into small blocks(subgraphs)
    Each block will have exact 1 input operation and 1 output operation.
    
    Besides, there is a parameter 'limit' to control the size of each block.
    """
    def __init__(self, graph: BaseGraph, topo_order: List[Operation]) -> None:
        self.graph = graph
        self.op_orders = topo_order
        self.depth = {}
        self.search_engine = SearchableGraph(self.graph)
        self.initialize_depth()

    def create_block(self, sp: Operation, ep: Operation) -> TrainableBlock:
        if sp == ep: return TrainableBlock(sp=sp, ep=ep, rps=[sp])
        rps = self.search_engine.opset_matching(
            sp_expr = lambda x: x == sp,
            rp_expr = lambda x, y: True,
            ep_expr = lambda x: x == ep,
            direction='down')
        rps = [(self.op_orders.index(op), op) for op in rps]
        rps = sorted(rps)
        return TrainableBlock(sp=sp, ep=ep, rps=[op for _, op in rps])

    def build(self, op: Operation, limit: int) -> TrainableBlock:
        """子图分割算法, 这个算法将从指定节点出发, 构造一个满足定义的子图结构 Solving best block from given
        operation.

        Block definition:(子图定义)
            A Block is a triple contains S, E, M,
                where S is the input node of block
                where E is the output node of block
                where M contains all nodes inside block

        Property:(子图性质)
            1. Minimal TrainableBlock start from p is {p, p, {p}},
               this block have only one node as both input and output.
            2. When S != E,
               E must on every path from S to graph output.
               S must on every path from graph input to E.
            3. M contains and only contains nodes on all paths from S to E.

        Lemma:(算法引理)
            1. 如果 s 的后继节点只有一个 e, 且 e 的输入只有一个，那么 {s, e, {s, e}} 构成满足定义的子图，从 s 寻找子图的任务可以递归由 e 完成
            2. 如果 s 的后继存在多个节点，则只存在两种情况:
                2.1 从 s 出发，最大子图即为 {s, s, {s}}。
                2.2 从 s 出发，可以构成子图 {s, e, R} (e!=s), 那么R中必须含有一个节点接收多个输入。（可用反证法证明，略）

        Algorithm:(算法)
            Build(s, d):
                如果区块长度大于所需，则返回现有内容
                从 s 出发，如果 s 的后继节点只有一个e，则判断e的输入节点个数：
                    1. 如果 e 是单输入节点，执行Build(e, d-1)，并将其结果与 {s, e, {s, e}} 合并
                    2. 如果 e 是多输入节点，算法立即停机，返回 {s, s, {s}}

                如果 s 的后继节点存在多个，找出距离 s 拓扑序最近的多输入的节点 k1，判断 s 到输出的路径是否能够被 k1 阻断
                    如果 k 成功阻断所有输出，执行Build(k1, d-1)，并将其结果与 {s, k1, F(s, k1)} 合并