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. (This resource utilization data is an estimate and not an exact value.)

Return type:

ResourceUtilization

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 (Union[TargetPlatformCapabilities, str]) – FrameworkQuantizationCapabilities 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)