sparseml.pytorch.utils package

Subpackages

• sparseml.pytorch.utils.quantization package
  • Submodules
  • sparseml.pytorch.utils.quantization.helpers module
  • sparseml.pytorch.utils.quantization.quantize_qat_export module
  • Module contents

Submodules

sparseml.pytorch.utils.benchmarker module

Benchmarking PyTorch models on a given device for given batch sizes

class sparseml.pytorch.utils.benchmarker.BatchBenchmarkResults(batch_size: int)

Bases: object

Container class for the results of a benchmark run for a given batch size. Contains convenience methods for calculating different metrics around the time to run each batch and the items.

Parameters

batch_size – the batch size the results are for

add(model_sec: float, e2e_sec: float, batch_size: int)

Add a new batch result

Parameters

• model_sec – the seconds it took to execute the model

• e2e_sec – the seconds it took to execute model and transfer to and from device

• batch_size – the size of the batch recorded

property batch_size

the batch size the results are for

Type

return

property e2e_batch_seconds

the average overall time to execute the batches through the model and the system. Includes model execution time as well as time to transfer the data to and from a device.

Type

return

property e2e_batch_timings

the overall timings in seconds for each batch to run through the model and the system. Includes model execution time as well as time to transfer the data to and from a device.

Type

return

property e2e_batches_per_second

inverse of e2e_batch_seconds

Type

return

property e2e_item_seconds

the batch averaged overall time to execute the batches through the model and the system (e2e_batch_seconds / batch_size). Includes model execution time as well as time to transfer the data to and from a device.

Type

return

property e2e_items_per_second

inverse of e2e_item_seconds

Type

return

property model_batch_seconds

the average time it took to execute the batches through the model. Does not include time for transferring data to and from device (if any)

Type

return

property model_batch_timings

the overall timings in seconds for each batch to run through the model. Does not include time for transferring data to and from device (if any)

Type

return

property model_batches_per_second

inverse of model_batch_seconds

Type

return

property model_item_seconds

the batch averaged time it took in seconds to execute one item through the model (model_batch_seconds / batch_size). Does not include time for transferring data to and from device (if any)

Type

return

property model_items_per_second

inverse of model_items_per_second

Type

return

class sparseml.pytorch.utils.benchmarker.ModuleBenchmarker(module: torch.nn.modules.module.Module)

Bases: object

Convenience class for benchmarking a model on a given device for given batches at a given precision.

Parameters

module – the module to benchmark

run_batches_on_device(batches: List[Any], device: str, full_precision: bool = True, test_size: int = 100, warmup_size: int = 10) → sparseml.pytorch.utils.benchmarker.BatchBenchmarkResults

Parameters

• batches – the batches to run through the model and benchmark, should all be of the same batch_size

• device – the device to run the model on, ex: cpu, cuda, cuda:0, cuda:0,1

• full_precision – True to run at float32, False to run at float16

• test_size – the number of batches to run and calculate timings over

• warmup_size – the number of batches to run before calculating timings

Returns

the batch results for benchmarking

sparseml.pytorch.utils.callbacks module

sparseml.pytorch.utils.callbacks.apply_one_hot_label_mapping(labels: torch.Tensor, class_names: Dict[Any, str])

sparseml.pytorch.utils.callbacks.cifar10_label_mapping(labels: torch.Tensor)

sparseml.pytorch.utils.callbacks.coco_mapping(labels: Iterable[torch.Tensor])

sparseml.pytorch.utils.callbacks.coco_yolo_2017_mapping(labels: Iterable[torch.Tensor])

sparseml.pytorch.utils.callbacks.imagenet_label_mapping(labels: torch.Tensor)

sparseml.pytorch.utils.callbacks.imagenette_label_mapping(labels: torch.Tensor)

sparseml.pytorch.utils.callbacks.mnist_label_mapping(labels: torch.Tensor)

sparseml.pytorch.utils.exporter module

Export PyTorch models to the local device

class sparseml.pytorch.utils.exporter.ModuleExporter(module: torch.nn.modules.module.Module, output_dir: str)

Bases: object

An exporter for exporting PyTorch modules into ONNX format as well as numpy arrays for the input and output tensors.

Parameters

• module – the module to export

• output_dir – the directory to export the module and extras to

export_onnx(sample_batch: Any, name: str = 'model.onnx', opset: int = 11, disable_bn_fusing: bool = True, convert_qat: bool = False, **export_kwargs)

Export an onnx file for the current module and for a sample batch. Sample batch used to feed through the model to freeze the graph for a particular execution.

Parameters

• sample_batch – the batch to export an onnx for, handles creating the static graph for onnx as well as setting dimensions

• name – name of the onnx file to save

• opset – onnx opset to use for exported model. Default is 11, if torch version is 1.2 or below, default is 9

• disable_bn_fusing – torch >= 1.7.0 only. Set True to disable batch norm fusing during torch export. Default and suggested setting is True. Batch norm fusing will change the exported parameter names as well as affect sensitivity analyses of the exported graph. Additionally, the DeepSparse inference engine, and other engines, perform batch norm fusing at model compilation.

• convert_qat – if True and quantization aware training is detected in the module being exported, the resulting QAT ONNX model will be converted to a fully quantized ONNX model using quantize_torch_qat_export. Default is False.

• export_kwargs – kwargs to be passed as is to the torch.onnx.export api call. Useful to pass in dyanmic_axes, input_names, output_names, etc. See more on the torch.onnx.export api spec in the PyTorch docs: https://pytorch.org/docs/stable/onnx.html

export_pytorch(optimizer: Optional[torch.optim.optimizer.Optimizer] = None, recipe: Optional[str] = None, epoch: Optional[int] = None, name: str = 'model.pth', use_zipfile_serialization_if_available: bool = True, include_modifiers: bool = False, export_entire_model: bool = False, arch_key: Optional[str] = None)

Export the pytorch state dicts into pth file within a pytorch framework directory.

Parameters

• optimizer – optional optimizer to export along with the module

• recipe – the recipe used to obtain the model

• epoch – optional epoch to export along with the module

• name – name of the pytorch file to save

• use_zipfile_serialization_if_available – for torch >= 1.6.0 only exports the Module’s state dict using the new zipfile serialization

• include_modifiers – if True, and a ScheduledOptimizer is provided as the optimizer, the associated ScheduledModifierManager and its Modifiers will be exported under the ‘manager’ key. Default is False

• export_entire_model – Exports entire file instead of state_dict

• arch_key – if provided, the arch_key will be saved in the checkpoint

export_samples(sample_batches: List[Any], sample_labels: Optional[List[Any]] = None, sample_originals: Optional[List[Any]] = None, exp_counter: int = 0)

Export a set list of sample batches as inputs and outputs through the model.

Parameters

• sample_batches – a list of the sample batches to feed through the module for saving inputs and outputs

• sample_labels – an optional list of sample labels that correspond to the the batches for saving

• exp_counter – the counter to start exporting the tensor files at

export_to_zoo(dataloader: torch.utils.data.dataloader.DataLoader, original_dataloader: Optional[torch.utils.data.dataloader.DataLoader] = None, shuffle: bool = False, max_samples: int = 20, data_split_cb: Optional[Callable[Any, Tuple[Any, Any]]] = None, label_mapping_cb: Optional[Callable[Any, Any]] = None, trace_script: bool = False, fail_on_torchscript_failure: bool = True, export_entire_model: bool = False)

Creates and exports all related content of module including sample data, onnx, pytorch and torchscript.

Parameters

• dataloader – DataLoader used to generate sample data

• original_dataloader – Optional dataloader to obtain the untransformed image.

• shuffle – Whether to shuffle sample data

• max_samples – Max number of sample data to create

• data_split_cb – Optional callback function to split data sample into a tuple (features,labels). If not provided will assume dataloader returns a tuple (features,labels).

• label_mapping_cb – Optional callback function to mapping dataset label to other formats.

• dataset_wrapper – Wrapper function for the dataset to add original data to each sample. If set to None will default to use the ‘iter_dataset_with_orig_wrapper’ function.

• trace_script – If true, creates torchscript via tracing. Otherwise, creates the torchscripe via scripting.

• fail_on_torchscript_failure – If true, fails if torchscript is unable to export model.

• export_entire_model – Exports entire file instead of state_dict

export_torchscript(name: str = 'model.pts', sample_batch: Optional[Any] = None)

Export the torchscript into a pts file within a framework directory. If a sample batch is provided, will create torchscript model in trace mode. Otherwise uses script to create torchscript.

Parameters

• name – name of the torchscript file to save

• sample_batch – If provided, will create torchscript model via tracing using the sample_batch

classmethod get_output_names(out: Any)

Get name of output tensors. Derived exporters specific to frameworks could override this method :param out: outputs of the model :return: list of names

sparseml.pytorch.utils.exporter.export_onnx(module: torch.nn.modules.module.Module, sample_batch: Any, file_path: str, opset: int = 11, disable_bn_fusing: bool = True, convert_qat: bool = False, dynamic_axes: Optional[Union[str, Dict[str, List[int]]]] = None, skip_input_quantize: bool = False, **export_kwargs)

Export an onnx file for the current module and for a sample batch. Sample batch used to feed through the model to freeze the graph for a particular execution.

Parameters

• module – torch Module object to export

• sample_batch – the batch to export an onnx for, handles creating the static graph for onnx as well as setting dimensions

• file_path – path to the onnx file to save

• opset – onnx opset to use for exported model. Default is 11, if torch version is 1.2 or below, default is 9

• disable_bn_fusing – torch >= 1.7.0 only. Set True to disable batch norm fusing during torch export. Default and suggested setting is True. Batch norm fusing will change the exported parameter names as well as affect sensitivity analyses of the exported graph. Additionally, the DeepSparse inference engine, and other engines, perform batch norm fusing at model compilation.

• convert_qat – if True and quantization aware training is detected in the module being exported, the resulting QAT ONNX model will be converted to a fully quantized ONNX model using quantize_torch_qat_export. Default is False.

• dynamic_axes – dictionary of input or output names to list of dimensions of those tensors that should be exported as dynamic. May input ‘batch’ to set the first dimension of all inputs and outputs to dynamic. Default is an empty dict

• skip_input_quantize – if True, the export flow will attempt to delete the first Quantize Linear Nodes(s) immediately after model input and set the model input type to UINT8. Default is False

• export_kwargs – kwargs to be passed as is to the torch.onnx.export api call. Useful to pass in dyanmic_axes, input_names, output_names, etc. See more on the torch.onnx.export api spec in the PyTorch docs: https://pytorch.org/docs/stable/onnx.html

sparseml.pytorch.utils.helpers module

Utility / helper functions

class sparseml.pytorch.utils.helpers.NamedLayerParam(layer_name, layer, param_name, param)

Bases: tuple

property layer

Alias for field number 1

property layer_name

Alias for field number 0

property param

Alias for field number 3

property param_name

Alias for field number 2

sparseml.pytorch.utils.helpers.any_str_or_regex_matches_param_name(param_name: str, name_or_regex_patterns: List[str]) → bool

Parameters

• param_name – The name of a parameter

• name_or_regex_patterns – List of full param names to match to the input or regex patterns to match with that should be prefixed with ‘re:’

Returns

True if any given str or regex pattern matches the given name

sparseml.pytorch.utils.helpers.default_device() → str

Returns

the device that should be defaulted to for the current setup. if multiple gpus are available then will return a string with all of them, else if single gpu available then will return cuda, else returns cpu

sparseml.pytorch.utils.helpers.device_of(inputs: Any)

sparseml.pytorch.utils.helpers.early_stop_data_loader(data_loader: torch.utils.data.dataloader.DataLoader, early_stop_steps: int)

An iterator that goes through the data_loader for yields and stops after early_stop_steps instead of the full loader

Parameters

• data_loader – the data loader to continually repeat

• early_stop_steps – if set, the number of steps to run and break out early instead of running all of the steps in the data loader, if < 1 then will run the full length

Returns

an iterable for the never ending data loader

sparseml.pytorch.utils.helpers.get_conv_layers(module: torch.nn.modules.module.Module) → Dict[str, torch.nn.modules.module.Module]

Parameters

module – the module to grab all conv layers for

Returns

a list of all the conv layers in the module

sparseml.pytorch.utils.helpers.get_layer(name: str, module: torch.nn.modules.module.Module) → torch.nn.modules.module.Module

Parameters

• name – the name of the layer to grab from the module

• module – the module containing the layer to grab

Returns

the module representing the layer in the module

sparseml.pytorch.utils.helpers.get_layer_param(param: str, layer: str, module: torch.nn.modules.module.Module) → torch.nn.parameter.Parameter

Parameters

• param – the name of the param to grab from the layer

• layer – the name of the layer to grab from the module

• module – the module containing the layer and the param

Returns

the param taken from the given layer in the module

sparseml.pytorch.utils.helpers.get_linear_layers(module: torch.nn.modules.module.Module) → Dict[str, torch.nn.modules.module.Module]

Parameters

module – the module to grab all linear layers for

Returns

a list of all linear layers in the module

sparseml.pytorch.utils.helpers.get_named_layers_and_params_by_regex(module: torch.nn.modules.module.Module, param_names: List[str], params_strict: bool = False) → List[sparseml.pytorch.utils.helpers.NamedLayerParam]

Parameters

• module – the module to get the matching layers and params from

• param_names – a list of names or regex patterns to match with full parameter paths. Regex patterns must be specified with the prefix ‘re:’

• params_strict – if True, this function will raise an exception if there a parameter is not found to match every name or regex in param_names

Returns

a list of NamedLayerParam tuples whose full parameter names in the given module match one of the given regex patterns or parameter names

sparseml.pytorch.utils.helpers.get_optim_groups_learning_rates(optim: torch.optim.optimizer.Optimizer) → List[float]

Parameters

optim – The optimizer to get the learning rates for

Returns

get a list of tuples corresponding to the learning rates for the param groups in the optimizer

sparseml.pytorch.utils.helpers.get_optim_learning_rate(optim: torch.optim.optimizer.Optimizer) → float

Parameters

optim – The optimizer to get the learning rate for

Returns

convenience function to get the first learning rate for any of the param groups in the optimizer

sparseml.pytorch.utils.helpers.get_prunable_layers(module: torch.nn.modules.module.Module) → List[Tuple[str, torch.nn.modules.module.Module]]

Parameters

module – the module to get the prunable layers from

Returns

a list containing the names and modules of the prunable layers (Linear, ConvNd)

sparseml.pytorch.utils.helpers.get_quantizable_layers(module: torch.nn.modules.module.Module) → List[Tuple[str, torch.nn.modules.module.Module]]

Parameters

module – the module to get the quantizable layers from

Returns

a list containing the names and modules of the quantizable layers (Linear, Conv2d, Conv3d)

sparseml.pytorch.utils.helpers.get_terminal_layers(module: torch.nn.modules.module.Module) → Dict[str, torch.nn.modules.module.Module]

Parameters

module – the module to grab all terminal layers for

Returns

a list of all of the terminal layers in a model (ie not containers; so convs, linears, activations, etc)

sparseml.pytorch.utils.helpers.infinite_data_loader(data_loader: torch.utils.data.dataloader.DataLoader, early_stop_steps: int = - 1, cache: bool = False)

A never ending data loader that will keep repeating the one passed in. Will additionally cache the data if requested.

Parameters

• data_loader – the data loader to continually repeat

• early_stop_steps – if set, the number of steps to run and break out early instead of running all of the steps in the data loader

• cache – True to cache the results in memory and return those on subsequent requests, False otherwise

Returns

an iterable for the never ending data loader

sparseml.pytorch.utils.helpers.mask_difference(old_mask: torch.Tensor, new_mask: torch.Tensor) → torch.Tensor

Parameters

• old_mask – the old mask to compare against for calculating the difference

• new_mask – the new mask to compare with for calculating the difference

Returns

a tensor representing the change from the old_mask to the new_mask specifically values returned as 1.0 are newly unmasked (0.0 => 1.0) values returned as -1.0 are newly masked (1.0 => 0.0) values returned as 0.0 had no change in (0.0 => 0.0 or 1.0 => 1.0)

sparseml.pytorch.utils.helpers.replace_layer(module: torch.nn.modules.module.Module, name: str, replace: torch.nn.modules.module.Module) → torch.nn.modules.module.Module

General function to replace a layer in a module with the given new one.

Parameters

• module – the module to replace the layer in

• name – the name of the layer to replace the activation for

• replace – the module to replace the layer with

Returns

the original layer that was replaced

sparseml.pytorch.utils.helpers.set_deterministic_seeds(seed: int = 0)

Manually seeds the numpy, random, and torch packages. Also sets torch.backends.cudnn.deterministic to True :param seed: the manual seed to use. Default is 0

sparseml.pytorch.utils.helpers.set_optim_learning_rate(optim: torch.optim.optimizer.Optimizer, value: float, groups: Optional[List[int]] = None)

Parameters

• optim – The optimizer to set the learning rate for

• value – the learning rate to set for the optimizer, will set all param groups in the optim to this value

sparseml.pytorch.utils.helpers.tensor_density(tens: torch.Tensor, dim: Union[None, int, Iterable[int]] = None) → torch.Tensor

Parameters

• tens – the tensor to calculate the density for

• dim – the dimension(s) to split the calculations over; ex, can split over batch, channels, or combos

Returns

the density of the input tens, ie the fraction of numbers that are non zero

sparseml.pytorch.utils.helpers.tensor_export(tensor: Union[torch.Tensor, Dict[str, torch.Tensor], Iterable[torch.Tensor]], export_dir: str, name: str, npz: bool = True) → str

Parameters

• tensor – tensor to export to a saved numpy array file

• export_dir – the directory to export the file in

• name – the name of the file, .npy will be appended to it

• npz – True to export as an npz file, False otherwise

Returns

the path of the numpy file the tensor was exported to

sparseml.pytorch.utils.helpers.tensor_list_sparsity(tensors: List[torch.Tensor]) → float

Parameters

tensors – the list of tensors to calculate the sparsity for

Returns

the total sparsity of all tensors in the list

sparseml.pytorch.utils.helpers.tensor_sample(tens: torch.Tensor, sample_size: int, dim: Union[None, int, List[int], Tuple[int, …]] = None) → torch.Tensor

Parameters

• tens – the tensor to grab samples from

• sample_size – the number of samples to grab overall if dim is not supplied or per each dim if it is

• dim – the dimension(s) to split the samples over; ex, can split over batch, channels, or combos

Returns

the sampled tensor

sparseml.pytorch.utils.helpers.tensor_sparsity(tens: torch.Tensor, dim: Union[None, int, List[int], Tuple[int, …]] = None) → torch.Tensor

Parameters

• tens – the tensor to calculate the sparsity for

• dim – the dimension(s) to split the calculations over; ex, can split over batch, channels, or combos

Returns

the sparsity of the input tens, ie the fraction of numbers that are zero

sparseml.pytorch.utils.helpers.tensors_batch_size(tensors: Union[torch.Tensor, Iterable[torch.Tensor], Dict[Any, torch.Tensor]])

Default function for getting the batch size from a tensor or collection of tensors. Returns the batch size (zeroth index for shape) of the first found tensor.

Supported use cases:

• single tensor

• Dictionary of single tensors

• Dictionary of iterable of tensors

• Dictionary of dictionary of tensors

• Iterable of single tensors

• Iterable of iterable of tensors

• Iterable of dictionary of tensors

Parameters

tensors – the tensor or collection of tensors to get a batch size from, taken from the first found tensor

Returns

the batch size (0th element of shape) of the first contained tensor in the data

sparseml.pytorch.utils.helpers.tensors_export(tensors: Union[torch.Tensor, Iterable[torch.Tensor]], export_dir: str, name_prefix: str, counter: int = 0, break_batch: bool = False) → List[str]

Parameters

• tensors – the tensors to export to a saved numpy array file

• export_dir – the directory to export the files in

• name_prefix – the prefix name for the tensors to save as, will append info about the position of the tensor in a list or dict in addition to the .npy file format

• counter – the current counter to save the tensor at

• break_batch – treat the tensor as a batch and break apart into multiple tensors

Returns

the exported paths

sparseml.pytorch.utils.helpers.tensors_module_forward(tensors: Union[torch.Tensor, Iterable[torch.Tensor], Mapping[Any, torch.Tensor]], module: torch.nn.modules.module.Module, check_feat_lab_inp: bool = True) → Any

Default function for calling into a model with data for a forward execution. Returns the model result. Note, if an iterable the features to be passed into the model are considered to be at index 0 and other indices are for labels.

Supported use cases: single tensor, iterable with first tensor taken as the features to pass into the model

Parameters

• tensors – the data to be passed into the model, if an iterable the features to be passed into the model are considered to be at index 0 and other indices are for labels

• module – the module to pass the data into

• check_feat_lab_inp – True to check if the incoming tensors looks like it’s made up of features and labels ie a tuple or list with 2 items (typical output from a data loader) and will call into the model with just the first element assuming it’s the features False to not check

Returns

the result of calling into the model for a forward pass

sparseml.pytorch.utils.helpers.tensors_to_device(tensors: Union[torch.Tensor, Iterable[torch.Tensor], Dict[Any, torch.Tensor]], device: str) → Union[torch.Tensor, Iterable[torch.Tensor], Dict[Any, torch.Tensor]]

Default function for putting a tensor or collection of tensors to the proper device. Returns the tensor references after being placed on the proper device.

Supported use cases:

• single tensor

• Dictionary of single tensors

• Dictionary of iterable of tensors

• Dictionary of dictionary of tensors

• Iterable of single tensors

• Iterable of iterable of tensors

• Iterable of dictionary of tensors

Parameters

• tensors – the tensors or collection of tensors to put onto a device

• device – the string representing the device to put the tensors on, ex: ‘cpu’, ‘cuda’, ‘cuda:1’

Returns

the tensors or collection of tensors after being placed on the device

sparseml.pytorch.utils.helpers.tensors_to_precision(tensors: Union[torch.Tensor, Iterable[torch.Tensor], Dict[Any, torch.Tensor]], full_precision: bool) → Union[torch.Tensor, Iterable[torch.Tensor], Dict[Any, torch.Tensor]]

Parameters

• tensors – the tensors to change the precision of

• full_precision – True for full precision (float 32) and False for half (float 16)

Returns

the tensors converted to the desired precision

sparseml.pytorch.utils.helpers.thin_model_from_checkpoint(model: torch.nn.modules.module.Module, state_dict: Dict[str, Any])

Updates any Linear/Conv/BN layers in the given model to match their respective shapes in the given state dict. Purpose of compatibility when loading weight for a model from a checkpoint of the same architecture but with potentially structured thinning applied. Note that this function has no guarantees on accuracy, will only resize model parameters for loading compatibility. All adjustments done in place

Parameters

• model – model to potentially adjust parameter shapes of

• state_dict – state dict to infer parameter shapes from

sparseml.pytorch.utils.helpers.torch_distributed_zero_first(local_rank: int)

Decorator to make all processes in distributed training wait for each local 0 ranked process to do something. :param local_rank: the local rank of this process

sparseml.pytorch.utils.logger module

Contains code for loggers that help visualize the information from each modifier

class sparseml.pytorch.utils.logger.BaseLogger(name: str, enabled: bool = True)

Bases: abc.ABC

Base class that all modifier loggers must implement.

Parameters

• name – name given to the logger, used for identification

• enabled – True to log, False otherwise

property enabled

True to log, False otherwise

Type

return

log_hyperparams(params: Dict[str, float]) → bool

Parameters

params – Each key-value pair in the dictionary is the name of the hyper parameter and it’s corresponding value.

Returns

True if logged, False otherwise.

log_scalar(tag: str, value: float, step: Optional[int] = None, wall_time: Union[None, float] = None, **kwargs) → bool

Parameters

• tag – identifying tag to log the value with

• value – value to save

• step – global step for when the value was taken

• wall_time – global wall time for when the value was taken

• kwargs – additional logging arguments to support Python and custom loggers

Returns

True if logged, False otherwise.

log_scalars(tag: str, values: Dict[str, float], step: Optional[int] = None, wall_time: Union[None, float] = None, **kwargs) → bool

Parameters

• tag – identifying tag to log the values with

• values – values to save

• step – global step for when the values were taken

• wall_time – global wall time for when the values were taken

• kwargs – additional logging arguments to support Python and custom loggers

Returns

True if logged, False otherwise.

log_string(tag: str, string: str, step: Optional[int] = None, wall_time: Union[None, float] = None, **kwargs) → bool

Parameters

• tag – identifying tag to log the values with

• values – values to save

• step – global step for when the values were taken

• wall_time – global wall time for when the values were taken

• kwargs – additional logging arguments to support Python and custom loggers

Returns

True if logged, False otherwise.

property name

name given to the logger, used for identification

Type

return

save(file_path: str, **kwargs) → bool

Parameters

• file_path – path to a file to be saved

• kwargs – additional arguments that a specific logger might use

Returns

True if saved, False otherwise

class sparseml.pytorch.utils.logger.LambdaLogger(lambda_func: Callable[[Optional[str], Optional[Union[float, str]], Optional[Dict[str, float]], Optional[int], Optional[float], Optional[int]], bool], name: str = 'lambda', enabled: bool = True)

Bases: sparseml.pytorch.utils.logger.BaseLogger

Logger that handles calling back to a lambda function with any logs.

Parameters

• lambda_func – the lambda function to call back into with any logs. The expected call sequence is (tag, value, values, step, wall_time) -> bool The return type is True if logged and False otherwise.

• name – name given to the logger, used for identification; defaults to lambda

• enabled – True to log, False otherwise

property lambda_func

the lambda function to call back into with any logs. The expected call sequence is (tag, value, values, step, wall_time)

Type

return

log_hyperparams(params: Dict, level: Optional[int] = None) → bool

Parameters

params – Each key-value pair in the dictionary is the name of the hyper parameter and it’s corresponding value.

Returns

True if logged, False otherwise.

log_scalar(tag: str, value: float, step: Optional[int] = None, wall_time: Union[None, float] = None, level: Optional[int] = None) → bool

Parameters

• tag – identifying tag to log the value with

• value – value to save

• step – global step for when the value was taken

• wall_time – global wall time for when the value was taken, defaults to time.time()

• kwargs – additional logging arguments to support Python and custom loggers

Returns

True if logged, False otherwise.

log_scalars(tag: str, values: Dict[str, float], step: Optional[int] = None, wall_time: Union[None, float] = None, level: Optional[int] = None) → bool

Parameters

• tag – identifying tag to log the values with

• values – values to save

• step – global step for when the values were taken

• wall_time – global wall time for when the values were taken, defaults to time.time()

• kwargs – additional logging arguments to support Python and custom loggers

Returns

True if logged, False otherwise.

class sparseml.pytorch.utils.logger.LoggerManager(loggers: Optional[List[sparseml.pytorch.utils.logger.BaseLogger]] = None, log_frequency: Optional[float] = 0.1, log_python: bool = True, name: str = 'manager')

Bases: abc.ABC

Wrapper around loggers that handles log scheduling and handing off logs to intended loggers.

Parameters

• loggers – list of loggers assigned to this manager

• log_frequency – number of epochs or fraction of epochs to wait between logs

static epoch_to_step(epoch, steps_per_epoch)

property log_frequency

number of epochs or fraction of epochs to wait between logs

Type

return

log_hyperparams(params: Dict, log_types: Union[str, List[str]] = '__ALL__', level: Optional[int] = None)

Parameters

params – Each key-value pair in the dictionary is the name of the hyper parameter and it’s corresponding value.

log_ready(epoch, last_log_epoch)

Check if there is a logger that is ready to accept a log

Parameters

• epoch – current epoch log is requested at

• last_log_epoch – last time a log was recorder for this object

Returns

True if a logger is ready to accept a log.

log_scalar(tag: str, value: float, step: Optional[int] = None, wall_time: Union[None, float] = None, log_types: Union[str, List[str]] = '__ALL__', level: Optional[int] = None)

Parameters

• tag – identifying tag to log the value with

• value – value to save

• step – global step for when the value was taken

• wall_time – global wall time for when the value was taken

• kwargs – additional logging arguments to support Python and custom loggers

Returns

True if logged, False otherwise.

log_scalars(tag: str, values: float, step: Optional[int] = None, wall_time: Union[None, float] = None, log_types: Union[str, List[str]] = '__ALL__', level: Optional[int] = None)

Parameters

• tag – identifying tag to log the values with

• values – values to save

• step – global step for when the values were taken

• wall_time – global wall time for when the values were taken

• kwargs – additional logging arguments to support Python and custom loggers

Returns

True if logged, False otherwise.

log_string(tag: str, string: str, step: Optional[int] = None, wall_time: Union[None, float] = None, log_types: Union[str, List[str]] = '__ALL__', level: Optional[int] = None)

Parameters

• tag – identifying tag to log the values with

• values – values to save

• step – global step for when the values were taken

• wall_time – global wall time for when the values were taken

• kwargs – additional logging arguments to support Python and custom loggers

Returns

True if logged, False otherwise.

property loggers

list of loggers assigned to this manager

Type

return

property name

name given to the logger, used for identification

Type

return

save(file_path: str, **kwargs)

Parameters

• file_path – path to a file to be saved

• kwargs – additional arguments that a specific logger might use

property wandb

wandb module if initialized

Type

return

class sparseml.pytorch.utils.logger.PythonLogger(logger: Optional[logging.Logger] = None, log_level: Optional[int] = None, name: str = 'python', enabled: bool = True)

Bases: sparseml.pytorch.utils.logger.LambdaLogger

Modifier logger that handles printing values into a python logger instance.

Parameters

• logger – a logger instance to log to, if None then will create it’s own

• log_level – default level to log any incoming data at on the logging.Logger instance when an explicit log level isn’t provided

• name – name given to the logger, used for identification; defaults to python

• enabled – True to log, False otherwise

log_string(tag: Optional[str], string: Optional[str], step: Optional[int], wall_time: Union[None, float] = None, level: Optional[int] = None) → bool

Parameters

• tag – identifying tag to log the values with

• string – string to log

• step – global step for when the values were taken

• wall_time – global wall time for when the values were taken, defaults to time.time()

• level – level to log at. Corresponds to default logging package levels

Returns

True if logged, False otherwise.

property logger

a logger instance to log to, if None then will create it’s own

Type

return

class sparseml.pytorch.utils.logger.SparsificationGroupLogger(lambda_func: Optional[Callable[[Optional[str], Optional[float], Optional[Dict[str, float]], Optional[int], Optional[float]], bool]] = None, python: Optional[Union[bool, logging.Logger]] = None, python_log_level: int = 20, tensorboard: Optional[Union[bool, str, torch.utils.tensorboard.writer.SummaryWriter]] = None, wandb_: Optional[Union[bool, Dict]] = None, name: str = 'sparsification', enabled: bool = True)

Bases: sparseml.pytorch.utils.logger.BaseLogger

Modifier logger that handles outputting values to other supported systems. Supported ones include:

• Python logging

• Tensorboard

• Weights and Biases

• Lambda callback

All are optional and can be bulk disabled and enabled by this root.

Parameters

• lambda_func – an optional lambda function to call back into with any logs. The expected call sequence is (tag, value, values, step, wall_time) -> bool The return type is True if logged and False otherwise.

• python – an optional argument for logging to a python logger. May be a logging.Logger instance to log to, True to create a logger instance, or non truthy to not log anything (False, None)

• python_log_level – if python, the level to log any incoming data at on the logging.Logger instance

• tensorboard – an optional argument for logging to a tensorboard writer. May be a SummaryWriter instance to log to, a string representing the directory to create a new SummaryWriter to log to, True to create a new SummaryWriter, or non truthy to not log anything (False, None)

• wandb – an optional argument for logging to wandb. May be a dictionary to pass to the init call for wandb, True to log to wandb (will not call init), or non truthy to not log anything (False, None)

• name – name given to the logger, used for identification; defaults to sparsification

• enabled – True to log, False otherwise

property enabled

True to log, False otherwise

Type

return

log_hyperparams(params: Dict, level: Optional[int] = None)

Parameters

params – Each key-value pair in the dictionary is the name of the hyper parameter and it’s corresponding value.

log_scalar(tag: str, value: float, step: Optional[int] = None, wall_time: Union[None, float] = None, level: Optional[int] = None)

Parameters

• tag – identifying tag to log the value with

• value – value to save

• step – global step for when the value was taken

• wall_time – global wall time for when the value was taken, defaults to time.time()

log_scalars(tag: str, values: Dict[str, float], step: Optional[int] = None, wall_time: Union[None, float] = None, level: Optional[int] = None)

Parameters

• tag – identifying tag to log the values with

• values – values to save

• step – global step for when the values were taken

• wall_time – global wall time for when the values were taken, defaults to time.time()

property loggers

the created logger sub instances for this logger

Type

return

class sparseml.pytorch.utils.logger.TensorBoardLogger(log_path: Optional[str] = None, writer: Optional[torch.utils.tensorboard.writer.SummaryWriter] = None, name: str = 'tensorboard', enabled: bool = True)

Bases: sparseml.pytorch.utils.logger.LambdaLogger

Modifier logger that handles outputting values into a TensorBoard log directory for viewing in TensorBoard.

Parameters

• log_path – the path to create a SummaryWriter at. writer must be None to use if not supplied (and writer is None), will create a TensorBoard dir in cwd

• writer – the writer to log results to, if none is given creates a new one at the log_path

• name – name given to the logger, used for identification; defaults to tensorboard

• enabled – True to log, False otherwise

property writer

the writer to log results to, if none is given creates a new one at the log_path

Type

return

class sparseml.pytorch.utils.logger.WANDBLogger(init_kwargs: Optional[Dict] = None, name: str = 'wandb', enabled: bool = True)

Bases: sparseml.pytorch.utils.logger.LambdaLogger

Modifier logger that handles outputting values to Weights and Biases.

Parameters

• init_kwargs – the args to call into wandb.init with; ex: wandb.init(**init_kwargs). If not supplied, then init will not be called

• name – name given to the logger, used for identification; defaults to wandb

• enabled – True to log, False otherwise

static available() → bool

Returns

True if wandb is available and installed, False, otherwise

save(file_path: str) → bool

Parameters

file_path – path to a file to be saved

sparseml.pytorch.utils.loss module

Code related to convenience functions for controlling the calculation of losses and metrics. Additionally adds in support for knowledge distillation

class sparseml.pytorch.utils.loss.Accuracy

Bases: torch.nn.modules.module.Module

Class for calculating the accuracy for a given prediction and the labels for comparison. Expects the inputs to be from a range of 0 to 1 and sets a crossing threshold at 0.5 the labels are similarly rounded.

static calculate(pred: torch.Tensor, lab: torch.Tensor)

Parameters

• pred – the models prediction to compare with

• lab – the labels for the data to compare to

Returns

the calculated accuracy

forward(pred: torch.Tensor, lab: torch.Tensor) → torch.Tensor

Parameters

• pred – the models prediction to compare with

• lab – the labels for the data to compare to

Returns

the calculated accuracy

training: bool

class sparseml.pytorch.utils.loss.BinaryCrossEntropyLossWrapper(extras: Optional[Dict] = None)

Bases: sparseml.pytorch.utils.loss.LossWrapper

Convenience class for doing binary cross entropy loss calculations, ie the default loss function is TF.binary_cross_entropy_with_logits.

Parameters

extras – extras representing other metrics that should be calculated in addition to the loss

class sparseml.pytorch.utils.loss.CrossEntropyLossWrapper(extras: Optional[Dict] = None)

Bases: sparseml.pytorch.utils.loss.LossWrapper

Convenience class for doing cross entropy loss calculations, ie the default loss function is TF.cross_entropy.

Parameters

extras – extras representing other metrics that should be calculated in addition to the loss

class sparseml.pytorch.utils.loss.InceptionCrossEntropyLossWrapper(extras: Optional[Dict] = None, aux_weight: float = 0.4)

Bases: sparseml.pytorch.utils.loss.LossWrapper

Loss wrapper for training an inception model that has an aux output with cross entropy.

Defines the loss in the following way: aux_weight * cross_entropy(aux_pred, lab) + cross_entropy(pred, lab)

Additionally adds cross_entropy into the extras.

Parameters

• extras – extras representing other metrics that should be calculated in addition to the loss

• aux_weight – the weight to use for the cross_entropy value calculated from the aux output

get_preds(data: Any, pred: Tuple[torch.Tensor, torch.Tensor, torch.Tensor], name: str) → Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]

Override get_preds for the inception training output. Specifically expects the pred from the model to be a three tensor tuple: (aux logits, logits, classes)

For the loss function returns a tuple containing (aux logits, logits), for all other extras returns the logits tensor

Parameters

• data – data from a data loader

• pred – the prediction from an inception model, expected to be a tuple containing (aux logits, logits, classes)

• name – the name of the loss function that is asking for the information for calculation

Returns

the predictions from the model for the loss function; a tuple containing (aux logits, logits), for all other extras returns the logits tensor

loss(preds: Tuple[torch.Tensor, torch.Tensor], labels: torch.Tensor)

Loss function for inception to combine the overall outputs from the model along with the the auxiliary loss from an earlier point in the model

Parameters

• preds – the predictions tuple containing [aux output, output]

• labels – the labels to compare to

Returns

the combined cross entropy value

class sparseml.pytorch.utils.loss.KDLossWrapper(loss_fn: Callable[[Any, Any], torch.Tensor], extras: Union[None, Dict[str, Callable]] = None, deconstruct_tensors: bool = True, kd_settings: Union[None, sparseml.pytorch.utils.loss.KDSettings] = None)

Bases: sparseml.pytorch.utils.loss.LossWrapper

Special case of the loss wrapper that allows knowledge distillation. Makes some assumptions specifically for image classification tasks, so may not work out of the box for everything.

Parameters

• loss_fn – the loss function to calculate on forward call of this object, accessible in the returned Dict at DEFAULT_LOSS_KEY

• extras – extras representing other metrics that should be calculated in addition to the loss

• deconstruct_tensors – True to break the tensors up into expected predictions and labels, False to pass the tensors as is to loss and extras

• kd_settings – the knowledge distillation settings that guide how to calculate the total loss

forward(data: Any, pred: Any) → Dict[str, torch.Tensor]

override to calculate the knowledge distillation loss if kd_settings is supplied and not None

Parameters

• data – the input data to the model, expected to contain the labels

• pred – the predicted output from the model

Returns

a dictionary containing all calculated losses and metrics with the loss from the loss_fn at DEFAULT_LOSS_KEY

get_inputs(data: Any, pred: Any, name: str) → Any

overridable function that is responsible for extracting the inputs to the model from the input data to the model and the output from the model

Parameters

• data – data from a data loader, expected to contain a tuple of (features, labels)

• pred – the predicted output from a model

• name – the name of the loss function that is asking for the information for calculation

Returns

the input data for the model

class sparseml.pytorch.utils.loss.KDSettings(teacher: torch.nn.modules.module.Module, temp_student: float = 5.0, temp_teacher: float = 5.0, weight: float = 0.5, contradict_hinton: bool = False)

Bases: object

properties class for settings for applying knowledge distillation as part of the loss calculation.

Parameters

• teacher – the teacher that provides targets for the student to learn from

• temp_student – temperature coefficient for the student

• temp_teacher – temperature coefficient for the teacher

• weight – the weight for how much of the kd loss to use in proportion with the original loss

• contradict_hinton – in hinton’s original paper they included T^2 as a scaling factor some implementations dropped this factor so contradicting hinton does not scale by T^2

property contradict_hinton

in hinton’s original paper they included T^2 as a scaling factor some implementations dropped this factor so contradicting hinton does not scale by T^2

Type

return

property teacher

the teacher that provides targets for the student to learn from

Type

return

property temp_student

temperature coefficient for the student

Type

return

property temp_teacher

temperature coefficient for the teacher

Type

return

property weight

the weight for how much of the kd loss to use in proportion with the original loss

Type

return

class sparseml.pytorch.utils.loss.LossWrapper(loss_fn: Callable[[Any, Any], torch.Tensor], extras: Union[None, Dict[str, Callable]] = None, deconstruct_tensors: bool = True)

Bases: object

Generic loss class for controlling how to feed inputs and compare with predictions for standard loss functions and metrics.

Parameters

• loss_fn – the loss function to calculate on forward call of this object, accessible in the returned Dict at DEFAULT_LOSS_KEY

• extras – extras representing other metrics that should be calculated in addition to the loss

• deconstruct_tensors – True to break the tensors up into expected predictions and labels, False to pass the tensors as is to loss and extras

property available_losses

a collection of all the loss and metrics keys available for this instance

Type

return

forward(data: Any, pred: Any) → Dict[str, torch.Tensor]

Parameters

• data – the input data to the model, expected to contain the labels

• pred – the predicted output from the model

Returns

a dictionary containing all calculated losses and metrics with the loss from the loss_fn at DEFAULT_LOSS_KEY

get_labels(data: Any, pred: Any, name: str) → Any

overridable function that is responsible for extracting the labels for the loss calculation from the input data to the model

Parameters

• data – data from a data loader, expected to contain a tuple of (features, labels)

• pred – the predicted output from a model

• name – the name of the loss function that is asking for the information for calculation

Returns

the label for the data

get_preds(data: Any, pred: Any, name: str) → Any

overridable function that is responsible for extracting the predictions from a model’s output

Parameters

• data – data from a data loader

• pred – the prediction from the model, if it is a tensor returns this, if it is an iterable returns first

• name – the name of the loss function that is asking for the information for calculation

Returns

the predictions from the model for the loss function

class sparseml.pytorch.utils.loss.SSDLossWrapper(extras: Optional[Dict] = None)

Bases: sparseml.pytorch.utils.loss.LossWrapper

Loss wrapper for SSD models. Implements the loss as the sum of:

1. Confidence Loss: All labels, with hard negative mining

2. Localization Loss: Only on positive labels

Parameters

extras – extras representing other metrics that should be calculated in addition to the loss

get_preds(data: Any, pred: Tuple[torch.Tensor, torch.Tensor], name: str) → Tuple[torch.Tensor, torch.Tensor]

Override get_preds for SSD model output.

Parameters

• data – data from a data loader

• pred – the prediction from an ssd model: two tensors representing object location and object label respectively

• name – the name of the loss function that is asking for the information for calculation

Returns

the predictions from the model without any changes

loss(preds: Tuple[torch.Tensor, torch.Tensor], labels: Tuple[torch.Tensor, torch.Tensor, torch.Tensor])

Calculates the loss for a multibox SSD output as the sum of the confidence and localization loss for the positive samples in the predictor with hard negative mining.

Parameters

• preds – the predictions tuple containing [predicted_boxes, predicted_lables].

• labels – the labels to compare to

Returns

the combined location and confidence losses

class sparseml.pytorch.utils.loss.TopKAccuracy(topk: int = 1)

Bases: torch.nn.modules.module.Module

Class for calculating the top k accuracy for a given prediction and the labels for comparison; ie the top1 or top5 accuracy. top1 is equivalent to the Accuracy class

Parameters

topk – the numbers of buckets the model is considered to be correct within

static calculate(pred: torch.Tensor, lab: torch.Tensor, topk: int)

Parameters

• pred – the models prediction to compare with

• lab – the labels for the data to compare to

• topk – the number of bins to be within for the correct label

Returns

the calculated topk accuracy

forward(pred: torch.Tensor, lab: torch.Tensor) → torch.Tensor

Parameters

• pred – the models prediction to compare with

• lab – the labels for the data to compare to

Returns

the calculated topk accuracy

training: bool

class sparseml.pytorch.utils.loss.YoloLossWrapper(extras: Optional[Dict] = None, anchor_groups: Optional[List[torch.Tensor]] = None)

Bases: sparseml.pytorch.utils.loss.LossWrapper

Loss wrapper for Yolo models. Implements the loss as a sum of class loss, objectness loss, and GIoU

Parameters

• extras – extras representing other metrics that should be calculated in addition to the loss

• anchor_groups – List of n,2 tensors of the Yolo model’s anchor points for each output group

forward(data: Any, pred: Any) → Dict[str, torch.Tensor]

Parameters

• data – the input data to the model, expected to contain the labels

• pred – the predicted output from the model

Returns

a dictionary containing all calculated losses (default, giou, object, and class) and metrics with the loss from the loss_fn at DEFAULT_LOSS_KEY

get_preds(data: Any, pred: List[torch.Tensor], name: str) → List[torch.Tensor]

Override get_preds for SSD model output.

Parameters

• data – data from a data loader

• pred – the prediction from an ssd model: two tensors representing object location and object label respectively

• name – the name of the loss function that is asking for the information for calculation

Returns

the predictions from the model without any changes

loss(preds: List[torch.Tensor], labels: Tuple[torch.Tensor, torch.Tensor])

Calculates the loss for a Yolo model output as the sum of the box, object, and class losses

Parameters

• preds – the predictions list containing objectness, class, and location values for each detector in the Yolo model.

• labels – the labels to compare to

Returns

the combined box, object, and class losses

sparseml.pytorch.utils.model module

Code related to interacting with a trained model such as saving, loading, etc

sparseml.pytorch.utils.model.device_to_name_ids(device: str) → Tuple[str, Union[None, List[int]]]

Split a device string into a device and ids

Parameters

device – the device string to push to; ex: cpu, cuda, cuda:0,1

Returns

a tuple containing the device string and devices

sparseml.pytorch.utils.model.is_parallel_model(model: torch.nn.modules.module.Module) → bool

Parameters

model – the model to test

Returns

True if the given model is wrapped as a DataPararallel or DistributedDataParallel Module. False otherwise

sparseml.pytorch.utils.model.load_epoch(path: str, map_location: Union[None, str] = 'cpu') → Optional[int]

sparseml.pytorch.utils.model.load_model(path: str, model: torch.nn.modules.module.Module, strict: bool = False, ignore_error_tensors: Optional[List[str]] = None, fix_data_parallel: bool = True)

Load the state dict into a model from a given file.

Parameters

• path – the path to the pth file to load the state dict from. May also be a SparseZoo stub path preceded by ‘zoo:’ with the optional ?recipe_type= argument. If given a recipe type, the base model weights for that recipe will be loaded.

• model – the model to load the state dict into

• strict – True to enforce that all tensors match between the model and the file; False otherwise

• ignore_error_tensors – names of tensors to ignore if they are not found in either the model or the file

• fix_data_parallel – fix the keys in the model state dict if they look like they came from DataParallel type setup (start with module.). This removes “module.” all keys

sparseml.pytorch.utils.model.load_optimizer(path: str, optimizer: torch.optim.optimizer.Optimizer, map_location: Union[None, str] = 'cpu')

Load the state dict into an optimizer from a given file.

Parameters

• path – the path to the pth file to load the state dict from

• optimizer – the optimizer to load the state dict into

• map_location – the location to map the values to when loading the

Returns

the epoch saved in the file, if any

sparseml.pytorch.utils.model.model_to_device(model: torch.nn.modules.module.Module, device: Union[str, int], ddp: bool = False) → Tuple[torch.nn.modules.module.Module, str, Union[None, List[int]]]

The model to push onto a device or multiple devices.

Parameters

• model – the model to push to a device

• device – the device string to push to; ex: cpu, cuda, cuda:0,1. For DDP, device should be the local_rank int value; ex: 0

• ddp – set True to wrap module as a DDP object. If True, device should be set to the local_rank int value. Default is False

Returns

a tuple containing the model on desired device(s), the device name, and the ids for the device

sparseml.pytorch.utils.model.parallelize_model(model: torch.nn.modules.module.Module, ids: Union[None, List[int]]) → torch.nn.modules.module.Module

Data parallelize a model across multiple devices

Parameters

• model – the model to parallelize across multiple devices

• ids – the ides of the devices to parallelize across

Returns

a parallelized model

sparseml.pytorch.utils.model.save_model(path: str, model: torch.nn.modules.module.Module, optimizer: Optional[torch.optim.optimizer.Optimizer] = None, recipe: Optional[str] = None, epoch: Optional[int] = None, use_zipfile_serialization_if_available: bool = True, include_modifiers: bool = False, arch_key: Optional[str] = None)

Save a model’s state dict into a file at the given path. Additionally can save an optimizer’s state and the current epoch.

Parameters

• path – the path to save the file the states to

• model – the model to save state for

• optimizer – the optimizer, if any, to save state for

• recipe – the recipe used to obtain the model

• epoch – the epoch to save

• use_zipfile_serialization_if_available – for torch >= 1.6.0 only exports the model’s state dict using the new zipfile serialization

• include_modifiers – if True, and a ScheduledOptimizer is provided as the optimizer, the associated ScheduledModifierManager and its Modifiers will be exported under the ‘manager’ key. Default is False

• arch_key – if provided, the arch_key will be saved in the checkpoint

sparseml.pytorch.utils.model.script_model(path: str, model: torch.nn.modules.module.Module)

Convenience function which scripts the provided module into a TorchScript script and saves to provied path.

Parameters

• path – path to save torchscript

• model – module to convert to torchscript

sparseml.pytorch.utils.model.trace_model(path: str, model: torch.nn.modules.module.Module, sample_batch: Any)

Convenience function which traces the provided module using the sample batch into a TorchScript script and saves to provied path.

Parameters

• path – path to save torchscript

• model – module to convert to TorchScript

• sample_batch – sample batch to trace module with

sparseml.pytorch.utils.module module

Code related to running a module through training and testing over a dataset. Allows reporting of progress and override functions and hooks.

class sparseml.pytorch.utils.module.ModuleDeviceContext(use_mixed_precision: bool = False, world_size: int = 1)

Bases: object

Simple class to define device settings or context to be used when running a Module

Parameters

• use_mixed_precision – set True to execute model using mixed precision with torch.cuda.amp. Default is False

• world_size – the world size (total number of devices) used when running the given module using DistributedDataParallel. Losses will be scaled by the world size. Default is 1.

static default_context()

Returns

A ModuleDeviceContext with default settings enabled

property use_mixed_precision

True if mixed precision with torch.cuda.amp should be used. False otherwise

Type

return

property world_size

the world size (total number of devices) used when running the given module using DistributedDataParallel. Losses will be scaled by the world size

Type

return

class sparseml.pytorch.utils.module.ModuleRunFuncs

Bases: object

Functions used as callables to calculate or perform necessary operations for running a model through training or testing.

Lifecycle:

- data batch size callback

- data to device callback

- batch start hook

- data model forward callback

- batch forward hook

- loss calculation

- batch loss hook

- model backward callback

- batch backward hook

- optimizer / gradient update

- batch end hook

property batch_size

return used to calculate the batch size of a given grouping of tensors. Expected to be called with the output from a data loader and then return an int representing the batch size.

copy(run_funcs)

Copy the functions from the current instance into a new instance

Parameters

run_funcs – the instance to copy the functions into

property model_backward

return used to call backward for a given model and the calculated losses. Expected to be called with the model and the output from the loss function as a dict mapping of names to tensors returns nothing

property model_forward

return used to propagate a given grouping of tensors through a model and return the result. Expected to be called with the model and the output from a data loader then return the result from the model forward pass.

property to_device

return used to place a given grouping of tensors onto the proper device. Expected to be called with the output from a data loader and the desired device as a string then return the grouping on the proper device.

class sparseml.pytorch.utils.module.ModuleRunHooks

Bases: object

Container for hooks that can be added to module runs like training and testing for different stages of running a batch through a model.

Lifecycle:

- data batch size callback

- data to device callback

- batch start hook

- data model forward callback

- batch forward hook

- loss calculation

- batch loss hook

- model backward callback

- batch backward hook

- optimizer / gradient update

- batch end hook

invoke_batch_backward(counter: int, step_count: int, batch_size: int, data: Any, pred: Any, losses: Dict[str, torch.Tensor])

invoke_batch_end(counter: int, step_count: int, batch_size: int, data: Any, pred: Any, losses: Dict[str, torch.Tensor])

invoke_batch_forward(counter: int, step_count: int, batch_size: int, data: Any, pred: Any)

invoke_batch_loss(counter: int, step_count: int, batch_size: int, data: Any, pred: Any, losses: Dict[str, torch.Tensor])

invoke_batch_start(counter: int, step_count: int, batch_size: int, data: Any)

register_batch_backward_hook(hook: Callable[[int, int, int, Any, Any, Dict[str, torch.Tensor]], None])

Called after calling backward on the loss for the batch with the following info: (counter, step_count, batch_size, data, pred, losses) where counter is passed in to the run (ex: epoch), step_count is the number of items run so far, batch_size is the number of elements fed in the batch, data is the data output from the loader, pred is the result from the model after the forward, losses are the resulting loss dictionary

Parameters

hook – the hook to add that is called into when reached in the batch process

Returns

a removable handle to remove the hook when desired

register_batch_end_hook(hook: Callable[[int, int, int, Any, Any, Dict[str, torch.Tensor]], None])

Called after all calculations are done for the batch with the following info: (counter, step_count, batch_size, data, pred, losses) where counter is passed in to the run (ex: epoch), step_count is the number of items run so far, batch_size is the number of elements fed in the batch, data is the data output from the loader, pred is the result from the model after the forward, losses are the resulting loss dictionary

Parameters

hook – the hook to add that is called into when reached in the batch process

Returns

a removable handle to remove the hook when desired

register_batch_forward_hook(hook: Callable[[int, int, int, Any, Any], None]) → torch.utils.hooks.RemovableHandle

Called after forward execution of a batch in the model with the following info: (counter, step_count, batch_size, data, pred) where counter is passed in to the run (ex: epoch), step_count is the number of items run so far, batch_size is the number of elements fed in the batch, data is the data output from the loader, pred is the result from the model after the forward

Parameters

hook – the hook to add that is called into when reached in the batch process

Returns

a removable handle to remove the hook when desired

register_batch_loss_hook(hook: Callable[[int, int, int, Any, Any, Dict[str, torch.Tensor]], None])

Called after loss calculation of the batch with the following info: (counter, step_count, batch_size, data, pred, losses) where counter is passed in to the run (ex: epoch), step_count is the number of items run so far, batch_size is the number of elements fed in the batch, data is the data output from the loader, pred is the result from the model after the forward, losses are the resulting loss dictionary

Parameters

hook – the hook to add that is called into when reached in the batch process

Returns

a removable handle to remove the hook when desired

register_batch_start_hook(hook: Callable[[int, int, int, Any], None]) → torch.utils.hooks.RemovableHandle

Called at the start of a batch with the following info: (counter, step_count, batch_size, data) where counter is passed in to the run (ex: epoch), step_count is the number of items run so far, batch_size is the number of elements fed in the batch, data is the data output from the loader

Parameters

hook – the hook to add that is called into when reached in the batch process

Returns

a removable handle to remove the hook when desired

class sparseml.pytorch.utils.module.ModuleRunResults

Bases: object

Class containing the results / losses from a model run for training or testing Keeps all result values as a dictionary and Tensor containing all values

append(losses: Dict[str, torch.Tensor], batch_size: int)

add new losses to the current stored results

Parameters

• losses – the losses to be added

• batch_size – the batch size the losses were run for

result(key: str) → List[torch.Tensor]

The result of a single loss function

Parameters

key – the name of the loss function to get the results for

Returns

a list of tensors containing all of the results for that loss

result_list_tensor(key: str) → torch.Tensor

Get the results as a list tensor where all items have been stacked into the first index of the tensor.

Parameters

key – the name of the loss function to get the results for

Returns

a tensor containing all of the tensors for that result

result_mean(key: str) → torch.Tensor

The mean result of a single loss function

Parameters

key – the name of the loss function to get the mean result for

Returns

a single tensor containing the average of all the results for that loss

result_std(key: str) → torch.Tensor

The standard deviation of the result for a single loss function

Parameters

key – the name of the loss function to get the standard deviation result for

Returns

a single tensor containing the standard deviation of all the results for that loss

property results

All of the stored results for the loss functions

Returns

a dictionary containing a mapping of name (str) to a list of tensors that were recorded for that loss

class sparseml.pytorch.utils.module.ModuleTester(module: torch.nn.modules.module.Module, device: str, loss: Union[sparseml.pytorch.utils.loss.LossWrapper, Callable[[Any, Any], torch.Tensor]], loggers: Optional[List[sparseml.pytorch.utils.logger.BaseLogger]] = None, log_name: str = 'Test', log_steps: int = 100, log_summary: bool = True, device_context: sparseml.pytorch.utils.module.ModuleDeviceContext = <sparseml.pytorch.utils.module.ModuleDeviceContext object>)

Bases: sparseml.pytorch.utils.module.ModuleRunner

Container for running a module through evaluation over a given data loader for specific settings.

Lifecycle:

- data batch size callback

- data to device callback

- batch start hook

- data model forward callback

- batch forward hook

- loss calculation

- batch loss hook

- batch end hook

Parameters

• module – the model to run evaluation for

• device – the default device to run evaluation on (where data will be copied to)

• loss – the loss functions callable used to calculate loss values after executing a forward pass

• loggers – list of loggers to log training results to

• log_name – the key to store all log files under

• log_steps – The number of steps (batches) to log at, ex 100 will log every 100 batches

• log_summary – True to log the final summary results after the run completes

• device_context – ModuleDeviceContext with settings to enable mixed precision using torch.cuda.amp or adjust losses when using DistributedDataParallel. Default settings do not use mixed precision or account for DDP. Will raise an exception if torch version does not support amp.

class sparseml.pytorch.utils.module.ModuleTrainer(module: torch.nn.modules.module.Module, device: str, loss: Union[sparseml.pytorch.utils.loss.LossWrapper, Callable[[Any, Any], torch.Tensor]], optimizer: torch.optim.optimizer.Optimizer, num_accumulated_batches: int = 1, optim_closure: Union[None, Callable] = None, loggers: Optional[List[sparseml.pytorch.utils.logger.BaseLogger]] = None, log_name: str = 'Train', log_steps: int = 100, log_summary: bool = True, device_context: sparseml.pytorch.utils.module.ModuleDeviceContext = <sparseml.pytorch.utils.module.ModuleDeviceContext object>)

Bases: sparseml.pytorch.utils.module.ModuleRunner

Container for running a module through training over a given data loader for specific settings.

Lifecycle:

- data batch size callback

- data to device callback

- batch start hook

- data model forward callback

- batch forward hook

- loss calculation

- batch loss hook

- model backward callback

- batch backward hook

- optimizer / gradient update

- batch end hook

Parameters

• module – the model to run training for

• device – the default device to run training on (where data will be copied to)

• loss – the loss functions callable used to calculate loss values after executing a forward pass

• optimizer – the optimizer used to apply gradient updates with

• num_accumulated_batches – number of batches to accumulate before updating the optimizer

• optim_closure – a closure passed into the optimizer on step

• loggers – list of loggers to log training results to

• log_name – the key to store all log files under

• log_steps – The number of steps (batches) to log at, ex 100 will log every 100 batches

• log_summary – True to log the final summary results after the run completes

• device_context – ModuleDeviceContext with settings to enable mixed precision using torch.cuda.amp or adjust losses when using DistributedDataParallel. Default settings do not use mixed precision or account for DDP. Will raise an exception if torch version does not support amp.

property num_accumulated_batches

number of batches to accumulate before updating the optimizer

Type

return

property optim_closure

a closure passed into the optimizer on step

Type

return

property optimizer

the optimizer used to apply gradient updates with

Type

return

sparseml.pytorch.utils.module.def_model_backward(losses: Dict[str, torch.Tensor], model: torch.nn.modules.module.Module, scaler: Optional[torch.cuda.amp.grad_scaler.GradScaler] = None)

Default function to perform a backwards pass for a model and the calculated losses Calls backwards for the DEFAULT_LOSS_KEY in losses Dict

Parameters

• model – the model to run the backward for

• losses – the losses dictionary containing named tensors, DEFAULT_LOSS_KEY is expected to exist and backwards is called on that

• scaler – GradScaler object for running in mixed precision with amp. If scaler is not None will call scaler.scale on the loss object. Default is None

sparseml.pytorch.utils.ssd_helpers module

Helper functions and classes for creating and training PyTorch SSD models

class sparseml.pytorch.utils.ssd_helpers.DefaultBoxes(image_size: int, feature_maps: List[int], steps: List[int], scales: List[int], aspect_ratios: List[List[int]], scale_xy: float = 0.1, scale_wh: float = 0.2)

Bases: object

Convenience class for creating, representing, encoding, and decoding default boxes

Parameters

• image_size – input image size

• feature_maps – list of feature map sizes

• steps – steps to use between boxes in a feature map

• scales – list of ranges of size scales to use for each feature map

• aspect_ratios – list of aspect ratios to construct boxes with

• scale_xy – parameter to scale box center by when encoding

• scale_wh – parameter to scale box dimensions by when encoding

as_ltrb() → torch.Tensor

Returns

The default boxes represented by this object in top left, top right pixel representation

as_xywh() → torch.Tensor

Returns

The default boxes represented by this object in center pixel, width, height representation

decode_output_batch(boxes: torch.Tensor, scores: torch.Tensor, score_threhsold: float = 0.01, iou_threshold: float = 0.45, max_detections: int = 200) → List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]

Decodes a batch detection model outputs from default box offsets and class scores to ltrb formatted bounding boxes, predicted labels, and scores for each image of the batch using non maximum suppression.

Parameters

• boxes – Encoded default-box offsets. Expected shape: batch_size,4,num_default_boxes

• scores – Class scores for each image, class, box combination. Expected shape: batch_size,num_classes,num_default_boxes

• score_threhsold – minimum softmax score to be considered a positive prediction. Default is 0.01 following the SSD paper

• iou_threshold – The minimum IoU between two boxes to be considered the same object in non maximum suppression

• max_detections – the maximum number of detections to keep per image. Default is 200

Returns

Detected object boudning boxes, predicted labels, and class score for each image in this batch

encode_image_box_labels(boxes: torch.Tensor, labels: torch.Tensor, threshold: float = 0.5) → Tuple[torch.Tensor, torch.Tensor]

Given the bounding box and image annotations for a single image with N objects will encode the box annotations as offsets to the default boxes and labels to the associated default boxes based on the annotation boxes and default boxes with an intersection over union (IoU) greater than the given threshold.

Parameters

• boxes – Bounding box annotations for objects in an image. Should have shape N,4 and be represented in ltrb format

• labels – Label annotations for N objects in an image.

• threshold – The minimum IoU bounding boxes and default boxes should share to be encoded

Returns

A tuple of the offset encoded bounding boxes and default box encoded labels

property num_default_boxes

the number of default boxes this object defines

Type

return

property scale_wh

parameter to scale box dimensions by when encoding

Type

return

property scale_xy

parameter to scale box center by when encoding

Type

return

class sparseml.pytorch.utils.ssd_helpers.MeanAveragePrecision(postprocessing_fn: Callable[Any, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]], iou_threshold: Union[float, Tuple[float, float]] = 0.5, iou_steps: float = 0.05)

Bases: object

Class for computing the mean average precision of an object detection model output. Inputs will be decoded by the provided post-processing function. Each batch update tracks the cumulative ground truth objects of each class, and the scores the model gives each class.

calculate_map object uses the aggregated results to find the mAP at the given threshold(s)

Parameters

• postprocessing_fn – function that takes in detection model output and returns post-processed tuple of predicted bounding boxes, classification labels, and scores

• iou_threshold – IoU thresholds to match predicted objects to ground truth objects. Can provide a single IoU or a tuple of two representing a range. mAP will be averaged over the range of values at each IoU

• iou_steps – the amount of IoU to shift between measurements between iou_threshold values

batch_forward(model_output: Tuple[torch.Tensor, torch.Tensor], ground_truth_annotations: List[Tuple[torch.Tensor, torch.Tensor]])

Decodes the model outputs using non maximum suppression, then stores the number of ground truth objects per class, true positives, and true negatives that can be used to calculate the overall mAP in the calculate_map function

Parameters

• model_output – the predictions tuple containing [predicted_boxes, predicted_labels] batch size should match length of ground_truth_annotations

• ground_truth_annotations – annotations from data loader to compare the batch results to, should be

calculate_map(num_recall_levels: int = 11) → Tuple[float, Dict[float, Dict[int, float]]]

Calculates mAP at the given threshold values based on the results stored in forward passes

Parameters

num_recall_levels – the number of recall levels to use between 0 and 1 inclusive. Defaults to 11, the VOC dataset standard

Returns

tuple of the overall mAP, and a dictionary that maps threshold level to class to average precision for that class

clear()

Resets the ground truth class count and results dictionaries

static get_recall_levels(num_recall_levels: int = 11)

Parameters

num_recall_levels – the number of recall levels to use between 0 and 1 inclusive. Defaults to 11, the VOC dataset standard

Returns

list of evenly spaced recall levels between 0 and 1 inclusive with num_recall_levels elements

sparseml.pytorch.utils.ssd_helpers.get_default_boxes_300(voc: bool = False) → sparseml.pytorch.utils.ssd_helpers.DefaultBoxes

Convenience function for generating DefaultBoxes object for standard SSD300 model

Parameters

voc – set True if default boxes should be made for VOC dataset. Will set scales to be slightly larger than for the default COCO dataset configuration

Returns

DefaultBoxes object implemented for standard SSD300 models

sparseml.pytorch.utils.ssd_helpers.ssd_random_crop(image: PIL.Image.Image, boxes: torch.Tensor, labels: torch.Tensor) → Tuple[PIL.Image.Image, torch.Tensor, torch.Tensor]

Performs one of the random SSD crops on a given image, bounding boxes, and labels as implemented in the original paper.

Chooses between following 3 conditions:

1. Preserve the original image

2. Random crop minimum IoU is among 0.1, 0.3, 0.5, 0.7, 0.9

3. Random crop

Adapted from: https://github.com/chauhan-utk/ssd.DomainAdaptation

Parameters

• image – the image to potentially crop

• boxes – a tensor of bounding boxes in ltrb format with shape n_boxes,4

• labels – a tensor of labels for each of the bounding boxes

Returns

the cropped image, boxes, and labels

sparseml.pytorch.utils.yolo_helpers module

Helper functions and classes for creating and training PyTorch Yolo models

class sparseml.pytorch.utils.yolo_helpers.YoloGrids(anchor_groups: Optional[List[torch.Tensor]] = None)

Bases: object

Helper class to compute and store Yolo output and anchor box grids

Parameters

anchor_groups – List of n,2 tensors of the Yolo model’s anchor points for each output group. Defaults to yolo_v3_anchor_groups

get_anchor_grid(group_idx: int) → torch.Tensor

Parameters

group_idx – Index of output group for this anchor grid

Returns

grid tensor of shape 1, num_anchors, 1, 1, 2

get_grid(size_x: int, size_y: int) → torch.Tensor

Parameters

• size_x – grid size x

• size_y – grid size y

Returns

Yolo output box grid for size x,y to be used for model output decoding. will have shape (1, 1, size_y, size_x, 2)

num_anchor_grids() → int

Returns

The number of anchor grids available (number of yolo model outputs)

sparseml.pytorch.utils.yolo_helpers.box_giou(boxes_a: torch.Tensor, boxes_b: torch.Tensor) → torch.Tensor

Parameters

• boxes_a – 4,N Tensor of xywh bounding boxes

• boxes_b – 4,N Tensor of xywh bounding boxes

Returns

Shape N Tensor of GIoU values between boxes in the input tensors

sparseml.pytorch.utils.yolo_helpers.build_targets(targets: torch.Tensor, anchors_groups: List[torch.Tensor], grid_shapes: List[torch.Tensor], iou_threshold: float = 0.2) → Tuple[List[torch.Tensor], List[torch.Tensor], List[torch.Tensor], List[torch.Tensor]]

Returns a representation of the image targets according to the given anchor groups and grid shapes.

Parameters

• targets – Yolo data targets tensor of shape n,6 with columns image number, class, center_x, center_y, width, height

• anchors_groups – List of n,2 Tensors of anchor point coordinates for each of the Yolo model’s detectors

• grid_shapes – List of n,2 Tensors of the Yolo models output grid shapes for a particular input shape

• iou_threshold – the minimum IoU value to consider an object box to match to an anchor point. Default is 0.2

Returns

sparseml.pytorch.utils.yolo_helpers.get_output_grid_shapes(outputs: List[torch.Tensor]) → List[torch.Tensor]

Parameters

outputs – List of Yolo model outputs

Returns

A list of the grid dimensions for each of the Yolo outputs

sparseml.pytorch.utils.yolo_helpers.postprocess_yolo(preds: List[torch.Tensor], input_shape: Iterable[int], yolo_grids: Optional[sparseml.pytorch.utils.yolo_helpers.YoloGrids] = None, confidence_threshold: float = 0.1, iou_threshold: float = 0.6, max_detections: int = 300) → List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]

Decode the outputs of a Yolo model and perform non maximum suppression on the predicted boxes.

Parameters

• preds – list of Yolo model output tensors

• input_shape – shape of input image to model. Default is [640, 640]

• yolo_grids – optional YoloGrids object for caching previously used grid shapes

• confidence_threshold – minimum confidence score for a prediction to be considered a detection. Default is 0.1

• iou_threshold – IoU threshold for non maximum suppression. Default is 0.6

• max_detections – maximum number of detections after nms. Default is 300

Returns

List of predicted bounding boxes (n,4), labels, and scores for each output in the batch

sparseml.pytorch.utils.yolo_helpers.yolo_v3_anchor_groups() → List[torch.Tensor]

Returns

List of the default anchor coordinate groups for Yolo V3 outputs

Module contents

Generic code used as utilities and helpers for PyTorch