nncf.torch.structures#

PyTorch-specific structure definitions for passing arguments into certain NNCF calls.

Classes#

QuantizationPrecisionInitArgs

Stores arguments for initialization of quantization's bitwidth.

AutoQPrecisionInitArgs

param data_loader:

'data_loader' - provides an iterable over the given dataset. Instance of

LeGRInitArgs

Stores arguments for learning global ranking in pruning algorithm.

DistributedCallbacksArgs

A pair of callbacks that is needed for distributed training of the model: wrapping model with wrapping_callback for

ExecutionParameters

Parameters that are necessary for distributed training of the model.
class nncf.torch.structures.QuantizationPrecisionInitArgs(criterion_fn, criterion, data_loader, device=None)[source]#

Bases: nncf.config.structures.NNCFExtraConfigStruct

Stores arguments for initialization of quantization’s bitwidth. Initialization is based on calculating a measure reflecting layers’ sensitivity to perturbations. The measure is calculated by estimation of average trace of Hessian for modules using the Hutchinson algorithm.

Parameters:

  • criterion_fn (Callable[[Any, Any, torch.nn.modules.loss._Loss], torch.Tensor]) – callable object, that implements calculation of loss by given outputs of the model, targets, and loss function. It’s not needed when the calculation of loss is just a direct call of the criterion with 2 arguments: outputs of model and targets. For all other specific cases, the callable object should be provided. E.g. for inception-v3, the losses for two outputs of the model are combined with different weight.

  • criterion (torch.nn.modules.loss._Loss) – loss function, instance of descendant of torch.nn.modules.loss._Loss,

  • data_loader (torch.utils.data.DataLoader) – ‘data_loader’ - provides an iterable over the given dataset. Instance of nncf.initialization.PTInitializingDataLoader; a regular ‘torch.utils.data.DataLoader’ may also be passed, but only in the simple case when it returns a tuple of (input, target) tensors.

  • device (str) –

Warning

The final quantizer setup of the created compressed model is dependent on the data provided by the data_loader. When using PyTorch’s DistributedDataParallel with precision initialization, make sure that each process in the distributed group receives the same data from the data_loader as the other processes, otherwise the create_compressed_model call may create different compressed model objects for each distributed process and the distributed training will fail.

Parameters:

  • device (str) – Device to perform initialization at. Either ‘cpu’, ‘cuda’, or None (default); if None, will use the device of the model’s parameters.

  • criterion_fn (Callable[[Any, Any, torch.nn.modules.loss._Loss], torch.Tensor]) –

  • criterion (torch.nn.modules.loss._Loss) –

  • data_loader (torch.utils.data.DataLoader) –

class nncf.torch.structures.AutoQPrecisionInitArgs(data_loader, eval_fn, nncf_config)[source]#

Bases: nncf.config.structures.NNCFExtraConfigStruct

Parameters:

  • data_loader (torch.utils.data.DataLoader) –

    ‘data_loader’ - provides an iterable over the given dataset. Instance of

    nncf.initialization.PTInitializingDataLoader; a regular ‘torch.utils.data.DataLoader’ may also be passed, but only in the simple case when it returns a tuple of (input, target) tensors.

    Warning

    The final quantizer setup of the created compressed model is dependent on the data provided by the data_loader. When using PyTorch’s DistributedDataParallel with precision initialization, make sure that each process in the distributed group receives the same data from the data_loader as the other processes, otherwise the create_compressed_model call may create different compressed model objects for each distributed process and the distributed training will fail.

  • eval_fn (Callable[[torch.nn.Module, torch.utils.data.DataLoader], float]) –

  • nncf_config (nncf.NNCFConfig) –

class nncf.torch.structures.LeGRInitArgs(train_loader, train_fn, val_loader, val_fn, train_optimizer, nncf_config)[source]#

Bases: nncf.config.structures.NNCFExtraConfigStruct

Stores arguments for learning global ranking in pruning algorithm.

Parameters:

  • train_loader (torch.utils.data.DataLoader) – provides an iterable over the given training (or initialising) dataset.

  • train_fn (Callable[[torch.utils.data.DataLoader, torch.nn.Module, torch.optim.Optimizer, nncf.api.compression.CompressionAlgorithmController, Optional[int]], type(None)]) – callable for training compressed model. Train model for one epoch or train_steps (if specified) by given args: [dataloader, model, optimizer, compression algorithm controller, train_steps number].

  • val_loader (torch.utils.data.DataLoader) – provides an iterable over the given validation dataset.

  • val_fn (Callable[[torch.nn.Module, torch.utils.data.DataLoader], Tuple[float, float]]) – callable to validate model, calculates pair of validation [acc, loss] by given model and dataloader.

  • train_optimizer (Optional[torch.optim.Optimizer]) – optional, optimizer for model training.

  • nncf_config (nncf.NNCFConfig) – NNCF config for compression.

class nncf.torch.structures.DistributedCallbacksArgs(wrapping_callback, unwrapping_callback)[source]#

Bases: nncf.config.structures.NNCFExtraConfigStruct

A pair of callbacks that is needed for distributed training of the model: wrapping model with wrapping_callback for distributed training, and after all training steps unwrapping model to the initial not-distributed state with unwrapping_callback.

Parameters:

  • wrapping_callback (Callable[[torch.nn.Module], torch.nn.Module]) – Callback that wraps the model for distributed training with any necessary structure (for example, torch.nn.DataParallel or any custom class), returns wrapped model ready for distributed training

  • unwrapping_callback (Callable[[torch.nn.Module], torch.nn.Module]) – Callback for unwrapping the model wrapped with wrapping_callback, returns original model

class nncf.torch.structures.ExecutionParameters(cpu_only, current_gpu)[source]#

Parameters that are necessary for distributed training of the model.

Parameters:

  • cpu_only (bool) – whether cpu-only mode is using for training

  • current_gpu (Optional[int]) – id of GPU that should be used for training (if only one of all is used)