afe.ir.transform.channel_scaling ================================ .. py:module:: afe.ir.transform.channel_scaling Attributes ---------- .. autoapisummary:: afe.ir.transform.channel_scaling.EPSILON afe.ir.transform.channel_scaling.SMOOTHQ_MINVAL afe.ir.transform.channel_scaling.EQUALIZATION_MINVAL afe.ir.transform.channel_scaling.EQUALIZATION_MAXVAL Classes ------- .. autoapisummary:: afe.ir.transform.channel_scaling.ScaleStatus afe.ir.transform.channel_scaling.PairingSet afe.ir.transform.channel_scaling.ResidualAddFamily Functions --------- .. autoapisummary:: afe.ir.transform.channel_scaling.get_pairing_lists afe.ir.transform.channel_scaling.get_pairings afe.ir.transform.channel_scaling.residual_pair_pass afe.ir.transform.channel_scaling.get_new_add_index afe.ir.transform.channel_scaling.get_add_connections afe.ir.transform.channel_scaling.get_equalization_scale afe.ir.transform.channel_scaling.parent_node_update afe.ir.transform.channel_scaling.child_node_update afe.ir.transform.channel_scaling.scale_concat_children afe.ir.transform.channel_scaling.scale_concat_parents afe.ir.transform.channel_scaling.find_largest_parent_set afe.ir.transform.channel_scaling.find_residual_scale afe.ir.transform.channel_scaling.pairings_update_pass Module Contents --------------- .. py:data:: EPSILON :value: 1e-30 .. py:data:: SMOOTHQ_MINVAL :value: 0.01 .. py:data:: EQUALIZATION_MINVAL :value: 1.0 .. py:data:: EQUALIZATION_MAXVAL :value: 20 .. py:class:: ScaleStatus Generic enumeration. Derive from this class to define new enumerations. .. py:attribute:: UNSCALED :value: 0 .. py:attribute:: TOSCALE :value: 1 .. py:attribute:: DONE :value: 2 .. py:class:: PairingSet A parent-children pairing. It contains the identified parents, children, and a status on whether the set has been scaled. .. py:attribute:: parents :type: list[afe.ir.node.AwesomeNode] .. py:attribute:: children :type: list[afe.ir.node.AwesomeNode] .. py:attribute:: status :type: ScaleStatus .. py:class:: ResidualAddFamily A collection of parents and children connected to an add-node. It contains the add-node, its parents and children, and a flag indicating if the node has been visited. .. py:attribute:: add_node :type: afe.ir.node.AwesomeNode .. py:attribute:: parents :type: list[afe.ir.node.AwesomeNode] .. py:attribute:: children :type: list[afe.ir.node.AwesomeNode] .. py:attribute:: visited :type: bool :value: False .. py:function:: get_pairing_lists(net: afe.ir.net.AwesomeNet) -> tuple[list[PairingSet], list[PairingSet], list[ResidualAddFamily]] Given an AwesomeNet, generate valid parent-children pairings that fit various patterns. The patterns are: 1. single-parent conv nodes with exclusively conv children. 2. Concat nodes with exclusively conv parents and children. 3. Residual add connections with exclusively conv parents and children. :param net: AwesomeNet to iterate through. :return: Tuple of Lists corresponding to each pattern. Each element of a list contains a pairing sets of parents and children of a pattern, and status for patterns 1 & 2, and a visited flag for pattern 3. .. py:function:: get_pairings(net: afe.ir.net.AwesomeNet) -> tuple[list[PairingSet], list[PairingSet], list[PairingSet]] Control function to find parent-children pairings to scale, and do necessary post-processing to generate the final and complete set of pairings. The goal of this pairing framework is to assemble self-contained parent-children 'pairing sets' that need to be scaled, such that each pairing set can be scaled without affecting the input activations to the set or the output activations from the set branching to other nodes in the network. Note that a node can appear in multiple pairing sets, but is always scaled at most once as a parent and at most once as a child. Thus, this framework only guarantees that activations going in and coming out of a set are invariant to scaling within a set, but does not guarantee that weight changes are self-contained to a single set. :param net: Network that is traversed to find node-pairs. :return: Tuple of Lists of parent-children pairing fitting each of the above listed patterns. .. py:function:: residual_pair_pass(add_node_family_list: list[ResidualAddFamily]) -> list[PairingSet] Iterate through the net to identify all the residual connection groupings, and identify all the parent and children nodes corresponding to each grouping. Each grouping will need to be scaled together with the same scale due to the design of the residual connection. A grouping is only valid if all its parents and children are convolution nodes (or pooling operators that lead to convolutions). Therefore, when a None is detected, the entire residual set is assembled and discarded from scaling. This algorithm is implemented by first identifying all the add operators in the network, and its immediate parents and children. It then scans each add-node's connections for further add-nodes, and recurses along the add-nodes connected together, pooling the corresponding parents and children together. After this, if even one parent/child is not a conv, the entire grouping is ineligible for scaling. :param add_node_family_list: List of pairing sets of add-nodes, its parents and children, and a visited flag for downstream use. :return: List of pairing sets, each containing the parents, children, and ScaleStatus of a residual grouping. .. py:function:: get_new_add_index(add_node_family_list: list[ResidualAddFamily], add_node_to_find: afe.ir.node.AwesomeNode) -> int Given a list of pairing sets containing add-nodes, identify the index of a given add-node. :param add_node_family_list: List of pairing sets, each containing a distinct add-node and its immediate parents and children. :param add_node_to_find: Desired add-node. :return: Index of desired add-node in list. .. py:function:: get_add_connections(add_node_family_list: list[ResidualAddFamily], add_node_idx: int) -> tuple[list[afe.ir.node.AwesomeNode], list[afe.ir.node.AwesomeNode]] For a given add-node, scan its parents, adding all convs to the list. If an add is detected, recurse to that add-node. If neither is detected, the grouping is invalid, so 'None' is appended to the list. Do this for the children of the add-node as well. :param add_node_family_list: List of pairing sets, each containing a distinct add-node and its immediate parents and children. :param add_node_idx: Index of the add-node we're currently scanning. :return: List of parents and children for a residual grouping. .. py:function:: get_equalization_scale(parent_node: afe.ir.node.AwesomeNode) -> tuple[numpy.ndarray, float, float] Compute a channel-wise scale given the node pair's parent node weights and output activation extrema. Based on the Same, Same, but Different paper: https://proceedings.mlr.press/v97/meller19a/meller19a.pdf. :param parent_node: The first node in the node-pair. :return: Channel-wise scales, maximum activation value, and new max after scaling. .. py:function:: parent_node_update(parent_node: afe.ir.node.AwesomeNode, scales: numpy.ndarray) Update a parent node's weights and bias given its scales. :param parent_node: AwesomeNode to update :param scales: Numpy array of scales to update along channels of parent node weight and bias. :return: None. .. py:function:: child_node_update(child_node: afe.ir.node.AwesomeNode, scales: numpy.ndarray) Update a child node's weights given its scales. :param child_node: AwesomeNode to update :param scales: Numpy array of scales to update along channels of parent node weight and bias. :return: None. .. py:function:: scale_concat_children(concat_node_pairing_list: list[PairingSet], parent_scale_dict: dict[afe.ir.node.AwesomeNode, numpy.ndarray]) Iterate through the list of identified concat node dicts. Whichever nodes are marked as 'to-do', concat the corresponding parent node scales, and rescale the children node weights. Mark the concat node's status as 'done' once children are scaled. :param concat_node_pairing_list: List of dicts of AwesomeNode parents and children that are connected by a concat node. :param parents_scale_dict: Dictionary of parent node names and scales :return: None .. py:function:: scale_concat_parents(concat_node_pairing: PairingSet) -> dict[afe.ir.node.AwesomeNode, numpy.ndarray] Iterate through a set of parent nodes. First find the original scales, then rescale so that all the max values are the same. Create a dictionary of the parent node names and their corresponding scales to scale the children nodes. :param concat_node_pairing: Pairing set for a given concat node, containing a list of parents, children, and status. We only use the parents list here. :return: Dictionary of parent node names and corresponding scales as numpy arrays. .. py:function:: find_largest_parent_set(concat_node_pairing_list: list[PairingSet], parents_list: list[afe.ir.node.AwesomeNode], idx_parents_list: int) -> int Iterate through the list of identified concat node pairing sets, and find the largest set of parents containing an initial subset of parent nodes. Update any concat set status as 'to-do' if the original set of parents are present. This function operates on the key asssumption that all parent sets are subsets of larger parent sets. :param concat_node_pairing_list: List of pairing sets of AwesomeNode parents and children that are connected by a concat node. :param parents_list: List of the initial set of parents for which we are trying to find supersets. :param idx_parents_list: Index of parents_list in the full concat_node_pairing_list :return: Index of largest set of parents .. py:function:: find_residual_scale(parent_nodes: list[afe.ir.node.AwesomeNode]) -> numpy.ndarray Iterate through the list of parent nodes and find the minimum of normalized scales across all parents. :param parent_nodes: List of AwesomeNode parents. :return: Residual scales. .. py:function:: pairings_update_pass(pairings_lists: tuple[list[PairingSet], list[PairingSet], list[PairingSet]]) Iterate through the lists of identified single-conv-parent pairings, concat pairings, and residual-add pairings, and scale them channel-wise based on channel_equalization. For conv_node_pairing_list, there can only be one parent, but the parent may have multiple children. The scale is found from the parent node, and is used to inversely scale the children nodes. For concat_node_pairing_list, iterate through the list of identified concat node pairing sets. For a given concat, the algorithm finds the scales for all its parents. It then rescales the scales to have the same max value. It then scales the parent nodes' weights, concats the scales together, and inversely scales the children nodes' weights. The algorithm makes the key assumption that for a parent set of a concat node, any other concat node that shares one of the parents will share all the parents (i.e., concats are structured in hierarchical manner, where parent sets are subsets of larger parent sets). This assumption is also made by the Same, Same, but Different paper's channel equalization implementation: https://github.com/icml2019/equalization. For residual_node_pairing_list, iterate through the list of identified residual node-pairs and scale them channel-wise. Do this by first finding the minimum of normalized scales across all parents. Then scale all parents and children by the same scale. :param pairings_lists: A tuple containing three lists of node pairings to scale: conv_node_pairing_list, concat_node_pairing_list, residual_node_pairing_list :return: Updates model weights in-place, returns nothing.