nodeclassifier

Node classification module for AutoCast.

This module provides classes for classifying ONNX nodes based on various rules to determine which nodes should be converted to lower precision and which should remain in high precision. It includes rules for handling node names, operation types, initializer ranges, and I/O value ranges.

Classes

`DisabledNodeNameRegexRule`	Rule for keeping nodes with matching names in high precision.
`DisabledOpTypes`	Rule for keeping nodes with specific operation types in high precision.
`IORangeRule`	Rule for keeping nodes with out-of-range inputs/outputs in high precision.
`InitializerRangeRule`	Rule for keeping nodes with out-of-range initializers in high precision.
`NodeClassifier`	Main class for classifying nodes into high and low precision groups.
`NodeRuleBase`	Base class for node classification rules.

class DisabledNodeNameRegexRule

Bases: NodeRuleBase

Rule for keeping nodes with matching names in high precision.

__init__(disabled_node_name_regex)

Initialize the rule.

Parameters:: disabled_node_name_regex – List of regex patterns for node names to keep in high precision.

class DisabledOpTypes

Bases: NodeRuleBase

Rule for keeping nodes with specific operation types in high precision.

__init__(op_types_to_exclude)

Initialize the rule.

Parameters:: op_types_to_exclude – List of operation types to keep in high precision.

class IORangeRule

Bases: NodeRuleBase

Rule for keeping nodes with out-of-range inputs/outputs in high precision.

__init__(data_max, reference_data, node_to_init_map)

Initialize the rule.

Parameters:

data_max – Maximum absolute value allowed for node I/O.
reference_data – Reference data for checking I/O ranges.
node_to_init_map – Mapping from node names to their initializers.

class InitializerRangeRule

Bases: NodeRuleBase

Rule for keeping nodes with out-of-range initializers in high precision.

__init__(init_max, node_to_init_map)

Initialize the rule.

Parameters:

init_max – Maximum absolute value allowed for initializers.
node_to_init_map – Mapping from node names to their initializers.

class NodeClassifier

Bases: object

Main class for classifying nodes into high and low precision groups.

__init__(model, node_to_init_map, nodes_to_exclude=None, op_types_to_exclude=[], custom_rule=None, data_max=1000, init_max=np.float16(65500.0))

Initialize the node classifier.

Parameters:

model – The ONNX model to classify nodes for.
node_to_init_map – Mapping from node names to their initializers.
nodes_to_exclude – List of regex patterns for node names to keep in high precision.
op_types_to_exclude – List of operation types to keep in high precision.
custom_rule – Optional custom classification rule.
data_max – Maximum absolute value allowed for node I/O.
init_max – Maximum absolute value allowed for initializers.

run(calibration_data=None)

Run node classification.

Parameters:: calibration_data – Optional input data for reference execution.
Returns:: Lists of node names (low_precision_nodes, high_precision_nodes).
Return type:: tuple

class NodeRuleBase

Bases: object

Base class for node classification rules.

This class defines the interface for rules that determine whether a node should be kept in high precision or converted to low precision.

check(node)

Check if a node should be skipped based on the rule.

Parameters:: node – The ONNX node to check.
Returns:: True if the node should be kept in high precision, False otherwise.
Return type:: bool