sparseml.pytorch.models.classification package

Submodules

sparseml.pytorch.models.classification.darknet module

DarkNet classification model for use as YOLO OD backbone

class sparseml.pytorch.models.classification.darknet.DarkNet(sec_settings: List[sparseml.pytorch.models.classification.darknet.DarkNetSectionSettings], num_classes: int, class_type: str)

Bases: torch.nn.modules.module.Module

DarkNet implementation

Parameters

• sec_settings – the settings for each section in the DarkNet model

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

as_classifier()

Sets this model to return output as an image classifier through a final FC layer

as_yolo_backbone(output_blocks: Optional[List[int]] = None)

Sets this model to output the given residual block indices as a backbone feature extractor for a detection model such as Yolo.

Parameters

output_blocks – indices of residual DarkNet blocks to output as backbone. Default is the final block output.

static create_section(settings: sparseml.pytorch.models.classification.darknet.DarkNetSectionSettings) → torch.nn.modules.container.Sequential

forward(inp: torch.Tensor) → Union[torch.Tensor, List[torch.Tensor]]

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class sparseml.pytorch.models.classification.darknet.DarkNetSectionSettings(num_blocks: int, in_channels: int, hidden_channels: int, out_channels: int, downsample_out_channels: Optional[int] = None)

Bases: object

Settings to describe how to put together a DarkNet based architecture using user supplied configurations.

Parameters

• num_blocks – the number of residual blocks to put in the section

• in_channels – the number of input channels to the section

• hidden_channels – the number of hidden channels in the residual blocks

• out_channels – the number of output channels for this sections residual blocks

• downsample_out_channels – number of output channels to apply to an additional convolution downsample layer. Setting to None will omit this layer. Default is None.

sparseml.pytorch.models.classification.darknet.darknet53(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single') → sparseml.pytorch.models.classification.darknet.DarkNet

DarkNet-53 implementation as described in the Yolo v3 paper; expected input shape is (B, 3, 256, 256)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created DarkNet Module

sparseml.pytorch.models.classification.efficientnet module

PyTorch EfficientNet implementation Further info can be found in the paper here.

class sparseml.pytorch.models.classification.efficientnet.EfficientNet(sec_settings: List[sparseml.pytorch.models.classification.efficientnet.EfficientNetSectionSettings], out_channels: int, num_classes: int, class_type: str, dropout: float)

Bases: torch.nn.modules.module.Module

EfficientNet implementation

Parameters

• sec_settings – the settings for each section in the vgg model

• out_channels – the number of output channels in the classifier before the fc

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• dropout – the amount of dropout to use while training

static create_section(settings: sparseml.pytorch.models.classification.efficientnet.EfficientNetSectionSettings) → torch.nn.modules.container.Sequential

forward(inp: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor]

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class sparseml.pytorch.models.classification.efficientnet.EfficientNetSectionSettings(num_blocks: int, in_channels: int, out_channels: int, kernel_size: int, expansion_ratio: int, stride: int, se_ratio: Optional[float], se_mod: bool)

Bases: object

Settings to describe how to put together an EfficientNet architecture using user supplied configurations.

Parameters

• num_blocks – the number of blocks to put in the section

• in_channels – the number of input channels to the section

• out_channels – the number of output channels from the section

• kernel_size – the kernel size of the depth-wise convolution

• expansion_ratio – (in_channels * expansion_ratio) is the number of input/output channels of the depth-wise convolution

• stride – the stride of the depth-wise convolution

• se_ratio – (in_channels * se_ratio) is the number of input channels for squeeze-excite

• se_mod – If true, moves squeeze-excite to the end of the block (after last 1x1)

sparseml.pytorch.models.classification.efficientnet.efficientnet_b0(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single', dropout: float = 0.2, se_mod: bool = False) → sparseml.pytorch.models.classification.efficientnet.EfficientNet

EfficientNet B0 implementation; expected input shape is (B, 3, 224, 224)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• dropout – the amount of dropout to use while training

• se_mod – If true, moves squeeze-excite to the end of the block (after last 1x1)

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created EfficientNet B0 Module

sparseml.pytorch.models.classification.efficientnet.efficientnet_b1(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single', dropout: float = 0.2, se_mod: bool = False) → sparseml.pytorch.models.classification.efficientnet.EfficientNet

EfficientNet B1 implementation; expected input shape is (B, 3, 240, 240)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• dropout – the amount of dropout to use while training

• se_mod – If true, moves squeeze-excite to the end of the block (after last 1x1)

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created EfficientNet B0 Module

sparseml.pytorch.models.classification.efficientnet.efficientnet_b2(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single', dropout: float = 0.3, se_mod: bool = False) → sparseml.pytorch.models.classification.efficientnet.EfficientNet

EfficientNet B2 implementation; expected input shape is (B, 3, 260, 260)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• dropout – the amount of dropout to use while training

• se_mod – If true, moves squeeze-excite to the end of the block (after last 1x1)

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created EfficientNet B0 Module

sparseml.pytorch.models.classification.efficientnet.efficientnet_b3(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single', dropout: float = 0.3, se_mod: bool = False) → sparseml.pytorch.models.classification.efficientnet.EfficientNet

EfficientNet B3 implementation; expected input shape is (B, 3, 300, 300)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• dropout – the amount of dropout to use while training

• se_mod – If true, moves squeeze-excite to the end of the block (after last 1x1)

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created EfficientNet B0 Module

sparseml.pytorch.models.classification.efficientnet.efficientnet_b4(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single', dropout: float = 0.4, se_mod: bool = False) → sparseml.pytorch.models.classification.efficientnet.EfficientNet

EfficientNet B4 implementation; expected input shape is (B, 3, 380, 380)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• dropout – the amount of dropout to use while training

• se_mod – If true, moves squeeze-excite to the end of the block (after last 1x1)

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created EfficientNet B0 Module

sparseml.pytorch.models.classification.efficientnet.efficientnet_b5(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single', dropout: float = 0.4, se_mod: bool = False) → sparseml.pytorch.models.classification.efficientnet.EfficientNet

EfficientNet B5 implementation; expected input shape is (B, 3, 456, 456)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• dropout – the amount of dropout to use while training

• se_mod – If true, moves squeeze-excite to the end of the block (after last 1x1)

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created EfficientNet B0 Module

sparseml.pytorch.models.classification.efficientnet.efficientnet_b6(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single', dropout: float = 0.5, se_mod: bool = False) → sparseml.pytorch.models.classification.efficientnet.EfficientNet

EfficientNet B6 implementation; expected input shape is (B, 3, 528, 528)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• dropout – the amount of dropout to use while training

• se_mod – If true, moves squeeze-excite to the end of the block (after last 1x1)

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created EfficientNet B0 Module

sparseml.pytorch.models.classification.efficientnet.efficientnet_b7(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single', dropout: float = 0.5, se_mod: bool = False) → sparseml.pytorch.models.classification.efficientnet.EfficientNet

EfficientNet B0 implementation; expected input shape is (B, 3, 600, 600)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• dropout – the amount of dropout to use while training

• se_mod – If true, moves squeeze-excite to the end of the block (after last 1x1)

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created EfficientNet B0 Module

sparseml.pytorch.models.classification.inception_v3 module

PyTorch Inception V3 implementations. Further info can be found in the paper here.

class sparseml.pytorch.models.classification.inception_v3.InceptionV3(num_classes: int, class_type: str, enable_aux: bool)

Bases: torch.nn.modules.module.Module

InceptionV3 implementation

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• enable_aux – True to enable the aux input for training, calculates aux logits from an earlier point in the network to enable smoother training as per the original paper

forward(x_tens: torch.Tensor) → Tuple[torch.Tensor, …]

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

sparseml.pytorch.models.classification.inception_v3.inception_v3(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single', enable_aux: bool = True) → sparseml.pytorch.models.classification.inception_v3.InceptionV3

Standard InceptionV3 implementation; expected input shape is (B, 3, 299, 299)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• enable_aux – True to enable the aux input for training, calculates aux logits from an earlier point in the network to enable smoother training as per the original paper

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created InceptionV3 Module

sparseml.pytorch.models.classification.mnist module

Simple PyTorch implementations for the MNIST dataset.

class sparseml.pytorch.models.classification.mnist.MnistNet(num_classes: int = 10, class_type: str = 'single')

Bases: torch.nn.modules.module.Module

A simple convolutional model created for the MNIST dataset

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

forward(inp: torch.Tensor)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

sparseml.pytorch.models.classification.mnist.mnist_net(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 10, class_type: str = 'single') → sparseml.pytorch.models.classification.mnist.MnistNet

MnistNet implementation; expected input shape is (B, 1, 28, 28)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created MnistNet Module

sparseml.pytorch.models.classification.mobilenet module

PyTorch MobileNet implementations. Further info can be found in the paper here.

class sparseml.pytorch.models.classification.mobilenet.MobileNet(sec_settings: List[sparseml.pytorch.models.classification.mobilenet.MobileNetSectionSettings], num_classes: int, class_type: str, dropout: Union[None, float] = None)

Bases: torch.nn.modules.module.Module

MobileNet implementation

Parameters

• sec_settings – the settings for each section in the MobileNet model

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• dropout – dropout level for input to FC layer; setting to None performs no dropout; default is None

static create_section(settings: sparseml.pytorch.models.classification.mobilenet.MobileNetSectionSettings) → torch.nn.modules.container.Sequential

forward(inp: torch.Tensor)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class sparseml.pytorch.models.classification.mobilenet.MobileNetSectionSettings(num_blocks: int, in_channels: int, out_channels: int, downsample: bool)

Bases: object

Settings to describe how to put together a MobileNet architecture using user supplied configurations.

Parameters

• num_blocks – the number of depthwise separable blocks to put in the section

• in_channels – the number of input channels to the section

• out_channels – the number of output channels from the section

• downsample – True to apply stride 2 for down sampling of the input, False otherwise

sparseml.pytorch.models.classification.mobilenet.han_mobilenet(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single') → sparseml.pytorch.models.classification.mobilenet.MobileNet

Standard MobileNet implementation with width=1.0; expected input shape is (B, 3, 224, 224)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created MobileNet Module

sparseml.pytorch.models.classification.mobilenet.mobilenet(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single') → sparseml.pytorch.models.classification.mobilenet.MobileNet

Standard MobileNet implementation with width=1.0; expected input shape is (B, 3, 224, 224)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created MobileNet Module

sparseml.pytorch.models.classification.mobilenet_v2 module

PyTorch MobileNet V2 implementations. Further info can be found in the paper here.

class sparseml.pytorch.models.classification.mobilenet_v2.MobilenetV2(sec_settings: List[sparseml.pytorch.models.classification.mobilenet_v2.MobilenetV2SectionSettings], num_classes: int, class_type: str)

Bases: torch.nn.modules.module.Module

Standard MobileNetV2 model https://arxiv.org/abs/1801.04381

static create_section(settings: sparseml.pytorch.models.classification.mobilenet_v2.MobilenetV2SectionSettings) → torch.nn.modules.container.Sequential

forward(inp: torch.Tensor)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class sparseml.pytorch.models.classification.mobilenet_v2.MobilenetV2SectionSettings(num_blocks: int, in_channels: int, out_channels: int, downsample: bool, exp_channels: Optional[int] = None, exp_ratio: float = 1.0, init_section: bool = False, width_mult: float = 1.0)

Bases: object

Settings to describe how to put together MobileNet V2 architecture using user supplied configurations.

Parameters

• num_blocks – the number of inverted bottleneck blocks to put in the section

• in_channels – the number of input channels to the section

• out_channels – the number of output channels from the section

• downsample – True to apply stride 2 for down sampling of the input, False otherwise

• exp_channels – number of channels to expand out to, if not supplied uses exp_ratio

• exp_ratio – the expansion ratio to use for the depthwise convolution

• init_section – True if it is the initial section, False otherwise

• width_mult – The width multiplier to apply to the channel sizes

sparseml.pytorch.models.classification.mobilenet_v2.mobilenet_v2(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single') → sparseml.pytorch.models.classification.mobilenet_v2.MobilenetV2

Standard MobileNet V2 implementation for a width multiplier; expected input shape is (B, 3, 224, 224)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created MobileNet Module

sparseml.pytorch.models.classification.mobilenet_v2.mobilenet_v2_width(width_mult: float, num_classes: int = 1000, class_type: str = 'single') → sparseml.pytorch.models.classification.mobilenet_v2.MobilenetV2

Standard MobileNet V2 implementation for a width multiplier; expected input shape is (B, 3, 224, 224)

Parameters

• width_mult – the width multiplier to apply

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

Returns

The created MobileNet Module

sparseml.pytorch.models.classification.resnet module

PyTorch ResNet, ResNet V2, ResNext implementations. Further info on ResNet can be found in the paper here. Further info on ResNet V2 can be found in the paper here. Further info on ResNext can be found in the paper here.

class sparseml.pytorch.models.classification.resnet.ResNet(sec_settings: List[sparseml.pytorch.models.classification.resnet.ResNetSectionSettings], num_classes: int, class_type: str)

Bases: torch.nn.modules.module.Module

ResNet, ResNet V2, ResNext implementations.

Parameters

• sec_settings – the settings for each section in the ResNet model

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

static create_section(settings: sparseml.pytorch.models.classification.resnet.ResNetSectionSettings) → torch.nn.modules.container.Sequential

forward(inp: torch.Tensor)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class sparseml.pytorch.models.classification.resnet.ResNetSectionSettings(num_blocks: int, in_channels: int, out_channels: int, downsample: bool, proj_channels: int = - 1, groups: int = 1, use_se: bool = False, version: int = 1)

Bases: object

Settings to describe how to put together a ResNet based architecture using user supplied configurations.

Parameters

• num_blocks – the number of blocks to put in the section (ie Basic or Bottleneck blocks)

• in_channels – the number of input channels to the section

• out_channels – the number of output channels from the section

• downsample – True to apply stride 2 for downsampling of the input, False otherwise

• proj_channels – The number of channels in the projection for a bottleneck block, if < 0 then uses basic

• groups – The number of groups to use for each 3x3 conv (ResNext)

• use_se – True to use squeeze excite, False otherwise

• version – 1 for original ResNet model, 2 for ResNet v2 model

sparseml.pytorch.models.classification.resnet.resnet101(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single') → sparseml.pytorch.models.classification.resnet.ResNet

Standard ResNet 101 implementation; expected input shape is (B, 3, 224, 224)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created ResNet Module

sparseml.pytorch.models.classification.resnet.resnet101_2xwidth(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single') → sparseml.pytorch.models.classification.resnet.ResNet

ResNet 101 implementation where channel sizes for 3x3 convolutions are doubled; expected input shape is (B, 3, 224, 224)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created ResNet Module

sparseml.pytorch.models.classification.resnet.resnet152(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single') → sparseml.pytorch.models.classification.resnet.ResNet

Standard ResNet 152 implementation; expected input shape is (B, 3, 224, 224)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created ResNet Module

sparseml.pytorch.models.classification.resnet.resnet18(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single') → sparseml.pytorch.models.classification.resnet.ResNet

Standard ResNet 18 implementation; expected input shape is (B, 3, 224, 224)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created ResNet Module

sparseml.pytorch.models.classification.resnet.resnet34(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single') → sparseml.pytorch.models.classification.resnet.ResNet

Standard ResNet 34 implementation; expected input shape is (B, 3, 224, 224)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created ResNet Module

sparseml.pytorch.models.classification.resnet.resnet50(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single') → sparseml.pytorch.models.classification.resnet.ResNet

Standard ResNet 50 implementation; expected input shape is (B, 3, 224, 224)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created ResNet Module

sparseml.pytorch.models.classification.resnet.resnet50_2xwidth(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single') → sparseml.pytorch.models.classification.resnet.ResNet

ResNet 50 implementation where channel sizes for 3x3 convolutions are doubled; expected input shape is (B, 3, 224, 224)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created ResNet Module

sparseml.pytorch.models.classification.resnet.resnetv2_101(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single') → sparseml.pytorch.models.classification.resnet.ResNet

Standard ResNet V2 101 implementation; expected input shape is (B, 3, 224, 224)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created ResNet Module

sparseml.pytorch.models.classification.resnet.resnetv2_152(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single') → sparseml.pytorch.models.classification.resnet.ResNet

Standard ResNet V2 152 implementation; expected input shape is (B, 3, 224, 224)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created ResNet Module

sparseml.pytorch.models.classification.resnet.resnetv2_18(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single') → sparseml.pytorch.models.classification.resnet.ResNet

Standard ResNet V2 18 implementation; expected input shape is (B, 3, 224, 224)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created ResNet Module

sparseml.pytorch.models.classification.resnet.resnetv2_34(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single') → sparseml.pytorch.models.classification.resnet.ResNet

Standard ResNet V2 34 implementation; expected input shape is (B, 3, 224, 224)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created ResNet Module

sparseml.pytorch.models.classification.resnet.resnetv2_50(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single') → sparseml.pytorch.models.classification.resnet.ResNet

Standard ResNet V2 50 implementation; expected input shape is (B, 3, 224, 224)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created ResNet Module

sparseml.pytorch.models.classification.resnet.resnext101(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single') → sparseml.pytorch.models.classification.resnet.ResNet

Standard ResNext 101 implementation; expected input shape is (B, 3, 224, 224)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created ResNet Module

sparseml.pytorch.models.classification.resnet.resnext152(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single') → sparseml.pytorch.models.classification.resnet.ResNet

Standard ResNext 152 implementation; expected input shape is (B, 3, 224, 224)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created ResNet Module

sparseml.pytorch.models.classification.resnet.resnext50(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single') → sparseml.pytorch.models.classification.resnet.ResNet

Standard ResNext 50 implementation; expected input shape is (B, 3, 224, 224)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created ResNet Module

sparseml.pytorch.models.classification.vgg module

PyTorch VGG implementations. Further info can be found in the paper here.

class sparseml.pytorch.models.classification.vgg.VGG(sec_settings: List[sparseml.pytorch.models.classification.vgg.VGGSectionSettings], num_classes: int, class_type: str)

Bases: torch.nn.modules.module.Module

VGG implementation

Parameters

• sec_settings – the settings for each section in the vgg model

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

static create_section(settings: sparseml.pytorch.models.classification.vgg.VGGSectionSettings) → torch.nn.modules.container.Sequential

forward(inp)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

sparseml.pytorch.models.classification.vgg.vgg11(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single') → sparseml.pytorch.models.classification.vgg.VGG

Standard VGG 11; expected input shape is (B, 3, 224, 224)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created MobileNet Module

sparseml.pytorch.models.classification.vgg.vgg11bn(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single') → sparseml.pytorch.models.classification.vgg.VGG

VGG 11 with batch norm added; expected input shape is (B, 3, 224, 224)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created MobileNet Module

sparseml.pytorch.models.classification.vgg.vgg13(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single') → sparseml.pytorch.models.classification.vgg.VGG

Standard VGG 13; expected input shape is (B, 3, 224, 224)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created MobileNet Module

sparseml.pytorch.models.classification.vgg.vgg13bn(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single') → sparseml.pytorch.models.classification.vgg.VGG

VGG 13 with batch norm added; expected input shape is (B, 3, 224, 224)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created MobileNet Module

sparseml.pytorch.models.classification.vgg.vgg16(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single') → sparseml.pytorch.models.classification.vgg.VGG

Standard VGG 16; expected input shape is (B, 3, 224, 224)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created MobileNet Module

sparseml.pytorch.models.classification.vgg.vgg16bn(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single') → sparseml.pytorch.models.classification.vgg.VGG

VGG 16 with batch norm added; expected input shape is (B, 3, 224, 224)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created MobileNet Module

sparseml.pytorch.models.classification.vgg.vgg19(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single') → sparseml.pytorch.models.classification.vgg.VGG

Standard VGG 19; expected input shape is (B, 3, 224, 224)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created MobileNet Module

sparseml.pytorch.models.classification.vgg.vgg19bn(pretrained_path: Optional[str] = None, pretrained: Union[bool, str] = False, pretrained_dataset: Optional[str] = None, load_strict: bool = True, ignore_error_tensors: Optional[List[str]] = None, num_classes: int = 1000, class_type: str = 'single') → sparseml.pytorch.models.classification.vgg.VGG

VGG 19 with batch norm added; expected input shape is (B, 3, 224, 224)

Parameters

• num_classes – the number of classes to classify

• class_type – one of [single, multi] to support multi class training; default single

• pretrained_path – A path to the pretrained weights to load, if provided will override the pretrained param. 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

• pretrained – True to load the default pretrained weights, a string to load a specific pretrained weight (ex: base, optim, optim-perf), or False to not load any pretrained weights

• pretrained_dataset – The dataset to load pretrained weights for (ex: imagenet, mnist, etc). If not supplied will default to the one preconfigured for the model.

• load_strict – True to raise an error on issues with state dict loading from pretrained_path or pretrained, False to ignore

• ignore_error_tensors – Tensors to ignore while checking the state dict for weights loaded from pretrained_path or pretrained

Returns

The created MobileNet Module

Module contents

Models related to image classification field in computer vision