afe.core.mixed_precision.annotation

Classes

IgnoreTensorPrecision

Functions

`remove_weight_quant`(→ onnx.ModelProto)	Removes weight quantizers in mct quantized model.
`update_precision`(→ onnx.ModelProto)	Updates precision of layers to promote in onnx annotated model.
`lookup_tensor_precision`(→ Union[None, ...)	Find the precision that mixed precision search has chosen for a tensor according to
`is_constant_tensor`(→ bool)	Return true if the tensor is an initializer or the output of a constant operator.
`get_redundant_holders`(→ List[str])	Find redundant Activation Holders from topologically sorted onnx nodes, only those Activation Holders
`is_initializer`(→ bool)	Check if a given string is an initializer.
`find_node_by_name`(→ Optional[onnx.NodeProto])	Find node given node name.
`find_node_by_output`(→ onnx.NodeProto)	Find node given node output name.
`find_node_by_input`(→ Tuple[onnx.NodeProto, int])	Find an operator node whose input is the specified input_name.
`remove_node`(→ None)	Remove node from model.
`replace_holders_with_annotations`(→ None)	Replace mct based activation holders with AFE specified annotations.
`set_sima_opset`(→ None)	Include sima opset.
`annotate_model`(→ None)	Main function to annotate model by updating activation holder precisions, remove redundant holders

Module Contents

afe.core.mixed_precision.annotation.remove_weight_quant(model: onnx.ModelProto) → onnx.ModelProto[source]

Removes weight quantizers in mct quantized model.

Parameters:: model – Onnx model with weight holders.
Returns:: model. Onnx model without weight holders.

afe.core.mixed_precision.annotation.update_precision(model: onnx.ModelProto, promotion_list: List[str]) → onnx.ModelProto[source]

Updates precision of layers to promote in onnx annotated model.

Parameters:

model – Onnx model with activation holders.
promotion_list – List of layer names to promote to 16 bit.

Returns:

model. Onnx model with activation holders with appropriate bit settings.

class afe.core.mixed_precision.annotation.IgnoreTensorPrecision[source]

afe.core.mixed_precision.annotation.lookup_tensor_precision(model: onnx.ModelProto, precs: Dict[str, int | None], tensor: str) → None | IgnoreTensorPrecision | int[source]

Find the precision that mixed precision search has chosen for a tensor according to attributes of Activation Holders.

Parameters:

model – Onnx model with activation holders.
precs – Precisions associated with nodes calculated so far. It holds at least the nodes that topologically precede the tensor.
tensor – Tensor to look up.

Returns:

Precision of the tensor.

afe.core.mixed_precision.annotation.is_constant_tensor(model: onnx.ModelProto, tensor: str) → bool[source]: Return true if the tensor is an initializer or the output of a constant operator.

afe.core.mixed_precision.annotation.get_redundant_holders(model: onnx.ModelProto) → List[str][source]

Find redundant Activation Holders from topologically sorted onnx nodes, only those Activation Holders that indicate a precision switch should be kept.

Parameters:: model – Onnx model with activation holders.
Returns:: remove_list. List of redundant activation holders to remove.

afe.core.mixed_precision.annotation.is_initializer(model: onnx.ModelProto, name: str) → bool[source]

Check if a given string is an initializer.

Parameters:

model – Onnx model with activation holders.
name – Name of initializer to search.

Returns:

True if initializer else False

afe.core.mixed_precision.annotation.find_node_by_name(model: onnx.ModelProto, node_name: str) → onnx.NodeProto | None[source]

Find node given node name.

Parameters:

model – Onnx model with activation holders.
node_name – Name of node to search.

Returns:

node. Node if it is found else None

afe.core.mixed_precision.annotation.find_node_by_output(model: onnx.ModelProto, output_name: str) → onnx.NodeProto[source]

Find node given node output name.

Parameters:

model – Onnx model with activation holders.
output_name – Name of node output to search.

Returns:

Node which output matches the output name. Raises an exception if node with specified output is not found.

afe.core.mixed_precision.annotation.find_node_by_input(model: onnx.ModelProto, input_name: str) → Tuple[onnx.NodeProto, int][source]

Find an operator node whose input is the specified input_name.

Parameters:

model – Loaded model in onnx.ModelProto representation.
input_name – Name of the input.

Returns:

Tuple of Node which input matches the input_name and input index. Raises an exception if node with specified input is not found.

afe.core.mixed_precision.annotation.remove_node(model: onnx.ModelProto, node_name: str) → None[source]

Remove node from model.

Parameters:

model – Onnx model with activation holders.
node_name – Name of the node to remove.

Returns:

None.

afe.core.mixed_precision.annotation.replace_holders_with_annotations(model: onnx.ModelProto) → None[source]

Replace mct based activation holders with AFE specified annotations.

Parameters:: model – Onnx model with activation holders.
Returns:: None.

afe.core.mixed_precision.annotation.set_sima_opset(model: onnx.ModelProto) → None[source]

Include sima opset.

Parameters:: model – Onnx model with activation holders.
Returns:: None.

afe.core.mixed_precision.annotation.annotate_model(promotion_list: List[str], annotated_onnx_filename: str) → None[source]

Main function to annotate model by updating activation holder precisions, remove redundant holders so that only the ones indicating a precision change are kept, replace holders with AFE specific annotations and save onnx model.

Parameters:

promotion_list – List of layer names to promote to 16 bit.
annotated_onnx_filename – Onnx file path used to load the original model and save the new model.

Returns:

None.