sparseml.pytorch.models.detection package

Submodules

sparseml.pytorch.models.detection.ssd module

Generic SSD Model framework

Information about SSD Networks can be found in the paper here.

class sparseml.pytorch.models.detection.ssd.SSD300(backbone: sparseml.pytorch.models.detection.ssd.SSDBackbone, num_classes: int = 91)

Bases: torch.nn.modules.module.Module

Single Shot Detector model that takes in a generic CNN backbone

Parameters

• backbone – SSDBackbone whose forward pass provides the feature extraction for an SSD model

• num_classes – number of target classes. Default is 81 to match the number of classes in the COCO detection dataset

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.detection.ssd.SSDBackbone

Bases: abc.ABC

Abstract class for representing backbone models for Single Shot Detectors

abstract get_feature_extractor() → torch.nn.modules.module.Module

Returns

A feature extrator module to be used for an SSD model

abstract property out_channels

a list of the sizes of the addtional out channels to be used with this backbone

Type

return

sparseml.pytorch.models.detection.ssd_lite module

Generic SSD Lite Model framework

Information about SSD Lite and mobilenet v2 can be found in the paper here.

class sparseml.pytorch.models.detection.ssd_lite.SSD300Lite(backbone: sparseml.pytorch.models.detection.ssd.SSDBackbone, backbone_early_output_idx: int, num_classes: int = 91)

Bases: torch.nn.modules.module.Module

Single Shot Detector model that takes in a CNN backbone and uses depthwise convolutions to perform fast object detection.

Parameters

• backbone – SSDBackbone whose forward pass provides the feature extraction for an SSD model

• backbone_early_output_idx – index of backbone Sequential object to be used as input to the first layer of the SSD head. All other layers use the entire backbone as input as per the SSDLite paper

• num_classes – number of target classes. Default is 81 to match the number of classes in the COCO detection dataset

forward(tens: 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

sparseml.pytorch.models.detection.ssd_mobilenet module

Implementations for SSD models with MobileNet backbones

class sparseml.pytorch.models.detection.ssd_mobilenet.SSD300MobileNetBackbone(version: Union[str, int] = '2', pretrained: Union[bool, str] = True, pretrained_path: Optional[str] = None)

Bases: sparseml.pytorch.models.detection.ssd.SSDBackbone

Class to provide the feature extractor and define the additional conv layers for an SSD300 model for various MobileNet architecture backbones

Parameters

• version – the MobileNet version to use for this backbone

• pretrained – True to load pretrained MobileNet weights; to load a specific version give a string with the name of the version (optim, optim-perf). Default is True

• pretrained_path – An optional model file path to load into the created model. Will override pretrained parameter

get_feature_extractor() → torch.nn.modules.module.Module

Returns

MobileNet feature extractor module to be used for an SSD model

property out_channels

The number of output channels that should be used for the additional conv layers with this backbone

Type

return

sparseml.pytorch.models.detection.ssd_mobilenet.ssd300lite_mobilenetv2(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 = 91, pretrained_backbone: Union[bool, str] = True, pretrained_path_backbone: str = None) → sparseml.pytorch.models.detection.ssd_lite.SSD300Lite

SSD 300 Lite with MobileNet V2 backbone; expected input shape is (B, 3, 300, 300)

Parameters

• num_classes – the number of classes of objects to classify

• pretrained_backbone – True to load pretrained MobileNet weights; to load a specific version give a string with the name of the version (optim, optim-perf). Default is True

• pretrained_path_backbone – An optional model file path to load into the created model’s backbone

• 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 SSD Lite MobileNet model

sparseml.pytorch.models.detection.ssd_resnet module

Implementations for SSD models with ResNet backbones

class sparseml.pytorch.models.detection.ssd_resnet.SSD300ResNetBackbone(sub_arch: Union[str, int], pretrained: Union[bool, str] = True, pretrained_path: Optional[str] = None)

Bases: sparseml.pytorch.models.detection.ssd.SSDBackbone

Class to provide the feature extractor and define the additional conv layers for an SSD300 model for various ResNet sub architecture backbones

Parameters

• sub_arch – the ResNet sub architecture to use for this backbone

• pretrained – True to load pretrained ResNet weights; to load a specific version give a string with the name of the version (optim, optim-perf). Default is True

• pretrained_path – An optional model file path to load into the created model. Will override pretrained parameter

get_feature_extractor() → torch.nn.modules.module.Module

Returns

ResNet feature extrator module to be used for an SSD model

property out_channels

The number of output channels that should be used for the additional conv layers with this backbone

Type

return

sparseml.pytorch.models.detection.ssd_resnet.ssd300_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 = 91, pretrained_backbone: Union[bool, str] = True, pretrained_path_backbone: str = None) → sparseml.pytorch.models.detection.ssd.SSD300

SSD 300 with ResNet 101 backbone; expected input shape is (B, 3, 300, 300)

Parameters

• num_classes – the number of classes of objects to classify

• pretrained_backbone – True to load pretrained ResNet weights; to load a specific version give a string with the name of the version (optim, optim-perf). Default is True

• pretrained_path_backbone – An optional model file path to load into the created model’s backbone

• 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 SSD ResNet model

sparseml.pytorch.models.detection.ssd_resnet.ssd300_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 = 91, pretrained_backbone: Union[bool, str] = True, pretrained_path_backbone: str = None) → sparseml.pytorch.models.detection.ssd.SSD300

SSD 300 with ResNet 152 backbone; expected input shape is (B, 3, 300, 300)

Parameters

• num_classes – the number of classes of objects to classify

• pretrained_backbone – True to load pretrained ResNet weights; to load a specific version give a string with the name of the version (optim, optim-perf). Default is True

• pretrained_path_backbone – An optional model file path to load into the created model’s backbone

• 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 SSD ResNet model

sparseml.pytorch.models.detection.ssd_resnet.ssd300_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 = 91, pretrained_backbone: Union[bool, str] = True, pretrained_path_backbone: str = None) → sparseml.pytorch.models.detection.ssd.SSD300

SSD 300 with ResNet 18 backbone; expected input shape is (B, 3, 300, 300)

Parameters

• num_classes – the number of classes of objects to classify

• pretrained_backbone – True to load pretrained ResNet weights; to load a specific version give a string with the name of the version (optim, optim-perf). Default is True

• pretrained_path_backbone – An optional model file path to load into the created model’s backbone

• 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 SSD ResNet model

sparseml.pytorch.models.detection.ssd_resnet.ssd300_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 = 91, pretrained_backbone: Union[bool, str] = True, pretrained_path_backbone: str = None) → sparseml.pytorch.models.detection.ssd.SSD300

SSD 300 with ResNet 34 backbone; expected input shape is (B, 3, 300, 300)

Parameters

• num_classes – the number of classes of objects to classify

• pretrained_backbone – True to load pretrained ResNet weights; to load a specific version give a string with the name of the version (optim, optim-perf). Default is True

• pretrained_path_backbone – An optional model file path to load into the created model’s backbone

• 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 SSD ResNet model

sparseml.pytorch.models.detection.ssd_resnet.ssd300_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 = 91, pretrained_backbone: Union[bool, str] = True, pretrained_path_backbone: str = None) → sparseml.pytorch.models.detection.ssd.SSD300

SSD 300 with ResNet 50 backbone; expected input shape is (B, 3, 300, 300)

Parameters

• num_classes – the number of classes of objects to classify

• pretrained_backbone – True to load pretrained ResNet weights; to load a specific version give a string with the name of the version (optim, optim-perf). Default is True

• pretrained_path_backbone – An optional model file path to load into the created model’s backbone

• 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 SSD ResNet model

sparseml.pytorch.models.detection.yolo_v3 module

PyTorch YoloV3 implementation.

class sparseml.pytorch.models.detection.yolo_v3.YoloV3(num_classes: int, backbone: torch.nn.modules.module.Module, backbone_out_channels: List[int], anchor_groups: List[torch.Tensor])

Bases: torch.nn.modules.module.Module

Yolo v3 implementation matching standard Yolo v3 SPP configuration

Parameters

• num_classes – the number of classes to classify objects with

• backbone – CNN backbone to this model

• backbone_out_channels – The number of output channels in each of the backbone’s outputs.

• anchor_groups – List of 3x2 Tensors of anchor point coordinates for each of this model’s detectors

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.detection.yolo_v3.yolo_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 = 80, pretrained_backbone: Union[bool, str] = False, pretrained_path_backbone: str = None) → sparseml.pytorch.models.detection.yolo_v3.YoloV3

Yolo-V3 model with standard DarkNet-53 backbone; expected input shape is (B, 3, 300, 300)

Parameters

• num_classes – the number of classes of objects to classify

• pretrained_backbone – True to load pretrained DarkNet weights; to load a specific version give a string with the name of the version (optim, optim-perf). Default is True

• pretrained_path_backbone – An optional model file path to load into the DarkNet backbone. Default is None

• 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 Yolo model

Module contents

Models related to object detection field in computer vision