Get Resource Utilization information for PyTorch Models

model_compression_toolkit.core.pytorch_resource_utilization_data(in_model, representative_data_gen, core_config=CoreConfig(), target_platform_capabilities=PYTORCH_DEFAULT_TPC)

Computes resource utilization data that can be used to calculate the desired target resource utilization for mixed-precision quantization. Builds the computation graph from the given model and target platform capabilities, and uses it to compute the resource utilization data.

Parameters:

  • in_model (Model) – PyTorch model to quantize.

  • representative_data_gen (Callable) – Dataset used for calibration.

  • core_config (CoreConfig) – CoreConfig containing parameters for quantization and mixed precision

  • target_platform_capabilities (TargetPlatformCapabilities) – TargetPlatformCapabilities to optimize the PyTorch model according to.

Returns:

A ResourceUtilization object with total weights parameters sum and max activation tensor.

Examples

Import a Pytorch model:

>>> from torchvision import models
>>> module = models.mobilenet_v2()

Create a random dataset generator:

>>> import numpy as np
>>> def repr_datagen(): yield [np.random.random((1, 3, 224, 224))]

Import mct and call for resource utilization data calculation:

>>> import model_compression_toolkit as mct
>>> ru_data = mct.core.pytorch_resource_utilization_data(module, repr_datagen)
Return type:

ResourceUtilization