afe.ir.transform.channel_scaling
Attributes
Classes
Generic enumeration. |
|
A parent-children pairing. It contains the identified parents, children, and |
|
A collection of parents and children connected to an add-node. It contains the |
Functions
|
Given an AwesomeNet, generate valid parent-children pairings that fit various |
|
Control function to find parent-children pairings to scale, and do necessary |
|
Iterate through the net to identify all the residual connection groupings, and identify all the parent and |
|
Given a list of pairing sets containing add-nodes, identify the index of a given add-node. |
|
For a given add-node, scan its parents, adding all convs to the list. If an add is detected, recurse |
|
Compute a channel-wise scale given the node pair's parent node |
|
Update a parent node's weights and bias given its scales. |
|
Update a child node's weights given its scales. |
|
Iterate through the list of identified concat node dicts. Whichever nodes are marked as 'to-do', |
|
Iterate through a set of parent nodes. First find the original scales, then rescale so that all the max values are |
|
Iterate through the list of identified concat node pairing sets, and find the largest set of parents containing an |
|
Iterate through the list of parent nodes and find the minimum of normalized scales across all |
|
Iterate through the lists of identified single-conv-parent pairings, concat pairings, and residual-add pairings, |
Module Contents
- class afe.ir.transform.channel_scaling.ScaleStatus[source]
Generic enumeration.
Derive from this class to define new enumerations.
- class afe.ir.transform.channel_scaling.PairingSet[source]
A parent-children pairing. It contains the identified parents, children, and a status on whether the set has been scaled.
- parents: list[afe.ir.node.AwesomeNode][source]
- children: list[afe.ir.node.AwesomeNode][source]
- status: ScaleStatus[source]
- class afe.ir.transform.channel_scaling.ResidualAddFamily[source]
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.
- add_node: afe.ir.node.AwesomeNode[source]
- parents: list[afe.ir.node.AwesomeNode][source]
- children: list[afe.ir.node.AwesomeNode][source]
- afe.ir.transform.channel_scaling.get_pairing_lists(net: afe.ir.net.AwesomeNet) tuple[list[PairingSet], list[PairingSet], list[ResidualAddFamily]][source]
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.
- Parameters:
net – AwesomeNet to iterate through.
- Returns:
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.
- afe.ir.transform.channel_scaling.get_pairings(net: afe.ir.net.AwesomeNet) tuple[list[PairingSet], list[PairingSet], list[PairingSet]][source]
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.
- Parameters:
net – Network that is traversed to find node-pairs.
- Returns:
Tuple of Lists of parent-children pairing fitting each of the above listed patterns.
- afe.ir.transform.channel_scaling.residual_pair_pass(add_node_family_list: list[ResidualAddFamily]) list[PairingSet][source]
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.
- Parameters:
add_node_family_list – List of pairing sets of add-nodes, its parents and children, and a visited flag for downstream use.
- Returns:
List of pairing sets, each containing the parents, children, and ScaleStatus of a residual grouping.
- afe.ir.transform.channel_scaling.get_new_add_index(add_node_family_list: list[ResidualAddFamily], add_node_to_find: afe.ir.node.AwesomeNode) int[source]
Given a list of pairing sets containing add-nodes, identify the index of a given add-node.
- Parameters:
add_node_family_list – List of pairing sets, each containing a distinct add-node and its immediate parents and children.
add_node_to_find – Desired add-node.
- Returns:
Index of desired add-node in list.
- afe.ir.transform.channel_scaling.get_add_connections(add_node_family_list: list[ResidualAddFamily], add_node_idx: int) tuple[list[afe.ir.node.AwesomeNode], list[afe.ir.node.AwesomeNode]][source]
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.
- Parameters:
add_node_family_list – List of pairing sets, each containing a distinct add-node and its immediate parents and children.
add_node_idx – Index of the add-node we’re currently scanning.
- Returns:
List of parents and children for a residual grouping.
- afe.ir.transform.channel_scaling.get_equalization_scale(parent_node: afe.ir.node.AwesomeNode) tuple[numpy.ndarray, float, float][source]
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.
- Parameters:
parent_node – The first node in the node-pair.
- Returns:
Channel-wise scales, maximum activation value, and new max after scaling.
- afe.ir.transform.channel_scaling.parent_node_update(parent_node: afe.ir.node.AwesomeNode, scales: numpy.ndarray)[source]
Update a parent node’s weights and bias given its scales.
- Parameters:
parent_node – AwesomeNode to update
scales – Numpy array of scales to update along channels of parent node weight and bias.
- Returns:
None.
- afe.ir.transform.channel_scaling.child_node_update(child_node: afe.ir.node.AwesomeNode, scales: numpy.ndarray)[source]
Update a child node’s weights given its scales.
- Parameters:
child_node – AwesomeNode to update
scales – Numpy array of scales to update along channels of parent node weight and bias.
- Returns:
None.
- afe.ir.transform.channel_scaling.scale_concat_children(concat_node_pairing_list: list[PairingSet], parent_scale_dict: dict[afe.ir.node.AwesomeNode, numpy.ndarray])[source]
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.
- Parameters:
concat_node_pairing_list – List of dicts of AwesomeNode parents and children that are connected by a concat node.
parents_scale_dict – Dictionary of parent node names and scales
- Returns:
None
- afe.ir.transform.channel_scaling.scale_concat_parents(concat_node_pairing: PairingSet) dict[afe.ir.node.AwesomeNode, numpy.ndarray][source]
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.
- Parameters:
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.
- Returns:
Dictionary of parent node names and corresponding scales as numpy arrays.
- afe.ir.transform.channel_scaling.find_largest_parent_set(concat_node_pairing_list: list[PairingSet], parents_list: list[afe.ir.node.AwesomeNode], idx_parents_list: int) int[source]
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.
- Parameters:
concat_node_pairing_list – List of pairing sets of AwesomeNode parents and children that are connected by a concat node.
parents_list – List of the initial set of parents for which we are trying to find supersets.
idx_parents_list – Index of parents_list in the full concat_node_pairing_list
- Returns:
Index of largest set of parents
- afe.ir.transform.channel_scaling.find_residual_scale(parent_nodes: list[afe.ir.node.AwesomeNode]) numpy.ndarray[source]
Iterate through the list of parent nodes and find the minimum of normalized scales across all parents.
- Parameters:
parent_nodes – List of AwesomeNode parents.
- Returns:
Residual scales.
- afe.ir.transform.channel_scaling.pairings_update_pass(pairings_lists: tuple[list[PairingSet], list[PairingSet], list[PairingSet]])[source]
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.
- Parameters:
pairings_lists – A tuple containing three lists of node pairings to scale: conv_node_pairing_list, concat_node_pairing_list, residual_node_pairing_list
- Returns:
Updates model weights in-place, returns nothing.