aimet_onnx.layer_output_utils

class aimet_onnx.layer_output_utils.LayerOutputUtil(model, dir_path, device=0)[source]

Implementation to capture and save outputs of intermediate layers of a model (fp32/quantsim)

Constructor - It initializes the utility classes that captures and saves layer-outputs

Parameters:

  • model (ModelProto) – ONNX model

  • dir_path (str) – Directory wherein layer-outputs will be saved

  • device (int) – CUDA device-id to be used

The following API can be used to Generate Layer Outputs

LayerOutputUtil.generate_layer_outputs(input_instance)[source]

This method captures output of every layer of a model & saves the inputs and corresponding layer-outputs to disk.

Parameters:

input_instance (Union[ndarray, List[ndarray], Tuple[ndarray]]) – Single input instance for which we want to obtain layer-outputs.

Returns:

None