sparseml.pytorch.utils package
Subpackages
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)[source] 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)[source] 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)[source] 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[source] - 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.exporter module
Export PyTorch models to the local device
-
class
sparseml.pytorch.utils.exporter.
ModuleExporter
(module: torch.nn.modules.module.Module, output_dir: str)[source] 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)[source] 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)[source] 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)[source] 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)[source] 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)[source] 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
-
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)[source] 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
-
property
-
sparseml.pytorch.utils.helpers.
any_str_or_regex_matches_param_name
(param_name: str, name_or_regex_patterns: List[str]) → bool[source] - 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[source] - 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.
early_stop_data_loader
(data_loader: torch.utils.data.dataloader.DataLoader, early_stop_steps: int)[source] 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][source] - 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[source] - 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[source] - 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][source] - 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][source] - 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][source] - 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[source] - 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]][source] - 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]][source] - 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][source] - 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)[source] 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[source] - 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[source] 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)[source] 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)[source] - 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[source] - 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[source] - 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[source] - 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[source] - 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[source] - 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]])[source] 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][source] - 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[source] 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]][source] 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]][source] - 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])[source] 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.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)[source] 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[source] - 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[source] - 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[source] - 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[source] - 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
-
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)[source] 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[source] - 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[source] - 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[source] - 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')[source] 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
-
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)[source] - 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)[source] 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)[source] - 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)[source] - 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)[source] - 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)[source] - 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)[source] 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[source] - 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)[source] 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)[source] - 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)[source] - 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)[source] - 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)[source] 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)[source] 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
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
[source] 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)[source] - 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[source] - Parameters
pred – the models prediction to compare with
lab – the labels for the data to compare to
- Returns
the calculated accuracy
-
training
: bool
-
static
-
class
sparseml.pytorch.utils.loss.
BinaryCrossEntropyLossWrapper
(extras: Optional[Dict] = None)[source] 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)[source] 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)[source] 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]][source] 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)[source] 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)[source] 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][source] 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[source] 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)[source] 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)[source] 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][source] - 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[source] 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[source] 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)[source] Bases:
sparseml.pytorch.utils.loss.LossWrapper
Loss wrapper for SSD models. Implements the loss as the sum of:
Confidence Loss: All labels, with hard negative mining
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][source] 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])[source] 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)[source] 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)[source] - 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[source] - 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)[source] 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][source] - 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][source] 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])[source] 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]]][source] 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[source] - 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][source]
-
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)[source] 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')[source] 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]]][source] 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[source] 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)[source] 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)[source] 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)[source] 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)[source] 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.
-
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
[source] 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)[source] 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.
-
property
-
class
sparseml.pytorch.utils.module.
ModuleRunHooks
[source] 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])[source]
-
invoke_batch_end
(counter: int, step_count: int, batch_size: int, data: Any, pred: Any, losses: Dict[str, torch.Tensor])[source]
-
invoke_batch_loss
(counter: int, step_count: int, batch_size: int, data: Any, pred: Any, losses: Dict[str, torch.Tensor])[source]
-
register_batch_backward_hook
(hook: Callable[[int, int, int, Any, Any, Dict[str, torch.Tensor]], None])[source] 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])[source] 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[source] 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])[source] 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[source] 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
[source] 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)[source] 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][source] 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[source] 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[source] 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[source] 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>)[source] 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>)[source] 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)[source] 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)[source] 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[source] - Returns
The default boxes represented by this object in top left, top right pixel representation
-
as_xywh
() → torch.Tensor[source] - 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]][source] 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][source] 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)[source] 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]])[source] 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]]][source] 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
-
sparseml.pytorch.utils.ssd_helpers.
get_default_boxes_300
(voc: bool = False) → sparseml.pytorch.utils.ssd_helpers.DefaultBoxes[source] 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][source] 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 image2. Random crop minimum IoU is among 0.1, 0.3, 0.5, 0.7, 0.93. Random cropAdapted 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)[source] 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[source] - Parameters
group_idx – Index of output group for this anchor grid
- Returns
grid tensor of shape 1, num_anchors, 1, 1, 2
-
sparseml.pytorch.utils.yolo_helpers.
box_giou
(boxes_a: torch.Tensor, boxes_b: torch.Tensor) → torch.Tensor[source] - 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]][source] 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][source] - 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]][source] 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
Module contents
Generic code used as utilities and helpers for PyTorch