Package Documentation

Check out the API section for quick reference.

sc2bench.common

sc2bench.common.config_util.overwrite_config(org_config, sub_config)

Overwrites a configuration.

Parameters:

• org_config (dict) – (nested) dictionary of configuration to be updated.

• sub_config (dict) – (nested) dictionary to be added to org_config.

sc2bench.models

sc2bench.models.backbone.register_backbone_class(cls)

Registers a backbone model (usually a classification model).

Parameters:

cls (class) – backbone model class to be registered

Returns:

registered backbone model class

Return type:

class

sc2bench.models.backbone.register_backbone_func(func)

Registers a function to build a backbone model (usually a classification model).

Parameters:

func (Callable) – function to build a backbone to be registered

Returns:

registered function

Return type:

Callable

class sc2bench.models.backbone.UpdatableBackbone(analyzer_configs=None)

A base, updatable R-CNN model.

Parameters:

analyzer_configs (list[dict] or None) – list of analysis configurations

update(**kwargs)

Updates compression-specific parameters like CompressAI models do.

This should be overridden by all subclasses.

get_aux_module(**kwargs)

Returns an auxiliary module to compute auxiliary loss if necessary like CompressAI models do.

This should be overridden by all subclasses.

sc2bench.models.backbone.check_if_updatable(model)

Checks if the given model is updatable.

Parameters:

model (nn.Module) – model

Returns:

True if the model is updatable, False otherwise

Return type:

bool

class sc2bench.models.backbone.FeatureExtractionBackbone(model, return_layer_dict, analyzer_configs, analyzes_after_compress=False, analyzable_layer_key=None)

A feature extraction-based backbone model.

Parameters:

• model (nn.Module) – model

• return_layer_dict (dict) – mapping from name of module to return its output to a specified key

• analyzer_configs (list[dict] or None) – list of analysis configurations

• analyzes_after_compress (bool) – run analysis with analyzer_configs if True

• analyzable_layer_key (str or None) – key of analyzable layer

check_if_updatable()

Checks if this module is updatable with respect to CompressAI modules.

Returns:

True if the model is updatable, False otherwise

Return type:

bool

update()

Updates compression-specific parameters like CompressAI models do.

This should be overridden by all subclasses.

get_aux_module(**kwargs)

Returns an auxiliary module to compute auxiliary loss if necessary like CompressAI models do.

This should be overridden by all subclasses.

class sc2bench.models.backbone.SplittableResNet(bottleneck_layer, resnet_model, inplanes=None, skips_avgpool=True, skips_fc=True, pre_transform_params=None, analysis_config=None)

ResNet/ResNeSt-based splittable image classification model containing neural encoder, entropy bottleneck, and decoder.

• Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun: “Deep Residual Learning for Image Recognition” @ CVPR 2016 (2016)

• Hang Zhang, Chongruo Wu, Zhongyue Zhang, Yi Zhu, Haibin Lin, Zhi Zhang, Yue Sun, Tong He, Jonas Mueller, R. Manmatha, Mu Li, Alexander Smola: “ResNeSt: Split-Attention Networks” @ CVPRW 2022 (2022)

• Yoshitomo Matsubara, Ruihan Yang, Marco Levorato, Stephan Mandt: “Supervised Compression for Resource-Constrained Edge Computing Systems” @ WACV 2022 (2022)

• Yoshitomo Matsubara, Ruihan Yang, Marco Levorato, Stephan Mandt: “SC2 Benchmark: Supervised Compression for Split Computing” @ TMLR (2023)

Parameters:

• bottleneck_layer (nn.Module) – high-level bottleneck layer that consists of encoder and decoder

• resnet_model (nn.Module) – ResNet model to be used as a base model

• inplanes (int or None) – ResNet model’s inplanes

• skips_avgpool (bool) – if True, skips avgpool (average pooling) after layer4

• skips_fc (bool) – if True, skips fc (fully-connected layer) after layer4

• pre_transform_params (dict or None) – pre-transform parameters

• analysis_config (dict or None) – analysis configuration

update()

Updates compression-specific parameters like CompressAI models do.

This should be overridden by all subclasses.

load_state_dict(state_dict, **kwargs)

Loads parameters for all the sub-modules except bottleneck_layer and then bottleneck_layer.

Parameters:

state_dict (dict) – dict containing parameters and persistent buffers

get_aux_module(**kwargs)

Returns an auxiliary module to compute auxiliary loss if necessary like CompressAI models do.

This should be overridden by all subclasses.

class sc2bench.models.backbone.SplittableRegNet(bottleneck_layer, regnet_model, inplanes=None, skips_head=True, pre_transform_params=None, analysis_config=None)

RegNet-based splittable image classification model containing neural encoder, entropy bottleneck, and decoder.

• Ilija Radosavovic, Raj Prateek Kosaraju, Ross Girshick, Kaiming He, Piotr Dollár: “Designing Network Design Spaces” @ CVPR 2020 (2020)

• Yoshitomo Matsubara, Ruihan Yang, Marco Levorato, Stephan Mandt: “SC2 Benchmark: Supervised Compression for Split Computing” @ TMLR (2023)

Parameters:

• bottleneck_layer (nn.Module) – high-level bottleneck layer that consists of encoder and decoder

• regnet_model (nn.Module) – RegNet model (timm-style) to be used as a base model

• inplanes (int or None) – mapping from name of module to return its output to a specified key

• skips_head (bool) – if True, skips fc (fully-connected layer) after layer4

• pre_transform_params (dict or None) – pre-transform parameters

• analysis_config (dict or None) – analysis configuration

update()

Updates compression-specific parameters like CompressAI models do.

This should be overridden by all subclasses.

load_state_dict(state_dict, **kwargs)

Loads parameters for all the sub-modules except bottleneck_layer and then bottleneck_layer.

Parameters:

state_dict (dict) – dict containing parameters and persistent buffers

get_aux_module(**kwargs)

Returns an auxiliary module to compute auxiliary loss if necessary like CompressAI models do.

This should be overridden by all subclasses.

class sc2bench.models.backbone.SplittableHybridViT(bottleneck_layer, hybrid_vit_model, num_pruned_stages=1, skips_head=True, pre_transform_params=None, analysis_config=None)

Hybrid ViT-based splittable image classification model containing neural encoder, entropy bottleneck, and decoder.

• Andreas Peter Steiner, Alexander Kolesnikov, Xiaohua Zhai, Ross Wightman, Jakob Uszkoreit, Lucas Beyer: “How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers” @ TMLR (2022)

• Yoshitomo Matsubara, Ruihan Yang, Marco Levorato, Stephan Mandt: “SC2 Benchmark: Supervised Compression for Split Computing” @ TMLR (2023)

Parameters:

• bottleneck_layer (nn.Module) – high-level bottleneck layer that consists of encoder and decoder

• hybrid_vit_model (nn.Module) – Hybrid Vision Transformer model (timm-style) to be used as a base model

• num_pruned_stages (int) – number of stages in the ResNet backbone of Hybrid ViT to be pruned

• skips_head (bool) – if True, skips classification head

• pre_transform_params (dict or None) – pre-transform parameters

• analysis_config (dict or None) – analysis configuration

update()

Updates compression-specific parameters like CompressAI models do.

This should be overridden by all subclasses.

load_state_dict(state_dict, **kwargs)

Loads parameters for all the sub-modules except bottleneck_layer and then bottleneck_layer.

Parameters:

state_dict (dict) – dict containing parameters and persistent buffers

get_aux_module(**kwargs)

Returns an auxiliary module to compute auxiliary loss if necessary like CompressAI models do.

This should be overridden by all subclasses.

sc2bench.models.backbone.splittable_resnet(bottleneck_config, resnet_name='resnet50', inplanes=None, skips_avgpool=True, skips_fc=True, pre_transform_params=None, analysis_config=None, org_model_ckpt_file_path_or_url=None, org_ckpt_strict=True, **resnet_kwargs)

Builds ResNet-based splittable image classification model containing neural encoder, entropy bottleneck, and decoder.

Parameters:

• bottleneck_config (dict) – bottleneck configuration

• resnet_name (str) – name of ResNet function in torchvision

• inplanes (int or None) – ResNet model’s inplanes

• skips_avgpool (bool) – if True, skips avgpool (average pooling) after layer4

• skips_fc (bool) – if True, skips fc (fully-connected layer) after layer4

• pre_transform_params (dict or None) – pre-transform parameters

• analysis_config (dict or None) – analysis configuration

• org_model_ckpt_file_path_or_url (str or None) – original ResNet model checkpoint file path or URL

• org_ckpt_strict (bool) – whether to strictly enforce that the keys in state_dict match the keys returned by original ResNet model’s state_dict() function

Returns:

splittable ResNet model

Return type:

SplittableResNet

sc2bench.models.backbone.splittable_resnest(bottleneck_config, resnest_name='resnest50d', inplanes=None, skips_avgpool=True, skips_fc=True, pre_transform_params=None, analysis_config=None, org_model_ckpt_file_path_or_url=None, org_ckpt_strict=True, **resnest_kwargs)

Builds ResNeSt-based splittable image classification model containing neural encoder, entropy bottleneck, and decoder.

Parameters:

• bottleneck_config (dict) – bottleneck configuration

• resnest_name (str) – name of ResNeSt function in timm

• inplanes (int or None) – ResNeSt model’s inplanes

• skips_avgpool (bool) – if True, skips avgpool (average pooling) after layer4

• skips_fc (bool) – if True, skips fc (fully-connected layer) after layer4

• pre_transform_params (dict or None) – pre-transform parameters

• analysis_config (dict or None) – analysis configuration

• org_model_ckpt_file_path_or_url (str or None) – original ResNeSt model checkpoint file path or URL

• org_ckpt_strict (bool) – whether to strictly enforce that the keys in state_dict match the keys returned by original ResNeSt model’s state_dict() function

Returns:

splittable ResNet model

Return type:

SplittableResNet

sc2bench.models.backbone.splittable_regnet(bottleneck_config, regnet_name='regnety_064', inplanes=None, skips_head=True, pre_transform_params=None, analysis_config=None, org_model_ckpt_file_path_or_url=None, org_ckpt_strict=True, **regnet_kwargs)

Builds RegNet-based splittable image classification model containing neural encoder, entropy bottleneck, and decoder.

Parameters:

• bottleneck_config (dict) – bottleneck configuration

• regnet_name (str) – name of RegNet function in timm

• inplanes (int or None) – mapping from name of module to return its output to a specified key

• skips_head (bool) – if True, skips fc (fully-connected layer) after layer4

• pre_transform_params (dict or None) – pre-transform parameters

• analysis_config (dict or None) – analysis configuration

• org_model_ckpt_file_path_or_url (str or None) – original RegNet model checkpoint file path or URL

• org_ckpt_strict (bool) – whether to strictly enforce that the keys in state_dict match the keys returned by original RegNet model’s state_dict() function

Returns:

splittable RegNet model

Return type:

SplittableRegNet

sc2bench.models.backbone.splittable_hybrid_vit(bottleneck_config, hybrid_vit_name='vit_small_r26_s32_224', num_pruned_stages=1, skips_head=True, pre_transform_params=None, analysis_config=None, org_model_ckpt_file_path_or_url=None, org_ckpt_strict=True, **hybrid_vit_kwargs)

Builds Hybrid ViT-based splittable image classification model containing neural encoder, entropy bottleneck, and decoder.

Parameters:

• bottleneck_config (dict) – bottleneck configuration

• hybrid_vit_name (str) – name of Hybrid ViT function in timm

• num_pruned_stages (int) – number of stages in the ResNet backbone of Hybrid ViT to be pruned

• skips_head (bool) – if True, skips classification head

• pre_transform_params (dict or None) – pre-transform parameters

• analysis_config (dict or None) – analysis configuration

• org_model_ckpt_file_path_or_url (str or None) – original Hybrid ViT model checkpoint file path or URL

• org_ckpt_strict (bool) – whether to strictly enforce that the keys in state_dict match the keys returned by original Hybrid ViT model’s state_dict() function

Returns:

splittable Hybrid ViT model

Return type:

SplittableHybridViT

sc2bench.models.backbone.get_backbone(cls_or_func_name, **kwargs)

Gets a backbone model.

Parameters:

• cls_or_func_name (str) – backbone class or function name

• kwargs (dict) – kwargs for the backbone class or function to build the backbone model

Returns:

backbone model

Return type:

nn.Module or None

sc2bench.models.layer.register_layer_class(cls)

Registers a layer class.

Parameters:

cls (class) – layer class to be registered

Returns:

registered layer class

Return type:

class

sc2bench.models.layer.register_layer_func(func)

Registers a function to build a layer module.

Parameters:

func (Callable) – function to build a layer module

Returns:

registered function

Return type:

Callable

class sc2bench.models.layer.SimpleBottleneck(encoder, decoder, compressor=None, decompressor=None)

Simple neural encoder-decoder that treats encoder’s output as bottleneck.

The forward path is encoder -> compressor (if provided) -> decompressor (if provided) -> decoder.

Parameters:

• encoder (nn.Module or None) – encoder

• decoder (nn.Module or None) – decoder

• encoder – module to compress the encoded data

• decoder – module to decompresse the compressed data

encode(x)

Encode the input data.

Parameters:

x (torch.Tensor) – input batch

Returns:

dict of encoded (and compressed if compressor is provided)

Return type:

dict

decode(z)

Decode the encoded data.

Parameters:

z (torch.Tensor) – encoded data

Returns:

decoded data

Return type:

torch.Tensor

update()

Shows a message that this module has no updatable parameters for entropy coding.

Dummy function to be compatible with other layers.

sc2bench.models.layer.larger_resnet_bottleneck(bottleneck_channel=12, bottleneck_idx=12, output_channel=256, compressor_transform_params=None, decompressor_transform_params=None)

Builds a bottleneck layer ResNet-based encoder and decoder (24 layers in total).

Compatible with ResNet-50, -101, and -152.

Yoshitomo Matsubara, Marco Levorato: “Neural Compression and Filtering for Edge-assisted Real-time Object Detection in Challenged Networks” @ ICPR 2020 (2021)

Parameters:

• bottleneck_channel – number of channels for the bottleneck point

• bottleneck_idx (int) – number of the first layers to be used as an encoder (the remaining layers are for decoder)

• output_channel (int) – number of output channels for decoder’s output

• compressor_transform_params (dict or None) – transform parameters for compressor

• decompressor_transform_params (dict or None) – transform parameters for decompressor

Returns:

bottleneck layer consisting of encoder and decoder

Return type:

SimpleBottleneck

class sc2bench.models.layer.EntropyBottleneckLayer(**kwargs)

An entropy bottleneck layer as a simple CompressionModel in compressai.

Johannes Ballé, David Minnen, Saurabh Singh, Sung Jin Hwang, Nick Johnston: “Variational Image Compression with a Scale Hyperprior” @ ICLR 2018 (2018)

Parameters:

kwargs (dict) – kwargs for CompressionModel in compressai

compress(x)

Compresses input data.

Parameters:

x (torch.Tensor) – input data

Returns:

entropy-coded compressed data (‘strings’ as key) and shape of the input data (‘shape’ as key)

Return type:

dict

decompress(strings, shape)

Dempresses compressed data.

Parameters:

• strings (list[str]) – entropy-coded compressed data

• shape (list[int]) – shape of the input data

Returns:

decompressed data

Return type:

torch.Tensor

update(force=False)

Updates compression-specific parameters like CompressAI models do.

Parameters:

force (bool) – if True, overwrites previous values

Returns:

True if one of the EntropyBottlenecks was updated

Return type:

bool

class sc2bench.models.layer.BaseBottleneck(entropy_bottleneck_channels)

An abstract class for entropy bottleneck-based layer.

Parameters:

entropy_bottleneck_channels (int) – number of entropy bottleneck channels

encode(*args, **kwargs)

Encodes data.

This should be overridden by all subclasses.

decode(*args, **kwargs)

Decodes encoded data.

This should be overridden by all subclasses.

update(force=False)

Updates compression-specific parameters like CompressAI models do.

Parameters:

force (bool) – if True, overwrites previous values

Returns:

True if one of the EntropyBottlenecks was updated

Return type:

bool

class sc2bench.models.layer.FPBasedResNetBottleneck(num_input_channels=3, num_bottleneck_channels=24, num_target_channels=256, encoder_channel_sizes=None, decoder_channel_sizes=None)

Factorized Prior(FP)-based encoder-decoder designed to create bottleneck for ResNet and variants.

• Johannes Ballé, David Minnen, Saurabh Singh, Sung Jin Hwang, Nick Johnston: “Variational Image Compression with a Scale Hyperprior” @ ICLR 2018 (2018)

• Yoshitomo Matsubara, Ruihan Yang, Marco Levorato, Stephan Mandt: “Supervised Compression for Resource-Constrained Edge Computing Systems” @ WACV 2022 (2022)

• Yoshitomo Matsubara, Ruihan Yang, Marco Levorato, Stephan Mandt: “SC2 Benchmark: Supervised Compression for Split Computing” @ TMLR (2023)

Parameters:

• num_input_channels (int) – number of input channels

• num_bottleneck_channels (int) – number of bottleneck channels

• num_target_channels (int) – number of output channels for decoder’s output

• encoder_channel_sizes (list[int] or None) – list of 4 numbers of channels for encoder

• decoder_channel_sizes (list[int] or None) – list of 4 numbers of channels for decoder

encode(x, **kwargs)

Encodes input data.

Parameters:

x (torch.Tensor) – input data

Returns:

entropy-coded compressed data (‘strings’ as key) and shape of the input data (‘shape’ as key)

Return type:

dict

decode(strings, shape)

Decodes encoded data.

Parameters:

• strings (list[str]) – entropy-coded compressed data

• shape (list[int]) – shape of the input data

Returns:

decompressed data

Return type:

torch.Tensor

class sc2bench.models.layer.SHPBasedResNetBottleneck(num_input_channels=3, num_latent_channels=16, num_bottleneck_channels=24, num_target_channels=256, h_a=None, h_s=None, g_a_channel_sizes=None, g_s_channel_sizes=None)

Scale Hyperprior(SHP)-based bottleneck for ResNet and variants.

• Johannes Ballé, David Minnen, Saurabh Singh, Sung Jin Hwang, Nick Johnston: “Variational Image Compression with a Scale Hyperprior” @ ICLR 2018 (2018)

• Yoshitomo Matsubara, Ruihan Yang, Marco Levorato, Stephan Mandt: “SC2 Benchmark: Supervised Compression for Split Computing” @ TMLR (2023)

Parameters:

• num_input_channels (int) – number of input channels

• num_latent_channels (int) – number of latent channels

• num_bottleneck_channels (int) – number of bottleneck channels

• num_target_channels (int) – number of output channels for decoder’s output

• h_a (nn.Module or None) – parametric transform \(h_a\)

• h_s (nn.Module or None) – parametric transform \(h_s\)

• g_a_channel_sizes (list[int] or None) – list of 4 numbers of channels for parametric transform \(g_a\)

• g_s_channel_sizes (list[int] or None) – list of 4 numbers of channels for parametric transform \(g_s\)

encode(x, **kwargs)

Encodes input data.

Parameters:

x (torch.Tensor) – input data

Returns:

entropy-coded compressed data (‘strings’ as key) and shape of the input data (‘shape’ as key)

Return type:

dict

decode(strings, shape)

Decodes encoded data.

Parameters:

• strings (list[str]) – entropy-coded compressed data

• shape (list[int]) – shape of the input data

Returns:

decompressed data

Return type:

torch.Tensor

update(scale_table=None, force=False)

Updates compression-specific parameters like CompressAI models do.

Parameters:

force (bool) – if True, overwrites previous values

Returns:

True if one of the EntropyBottlenecks was updated

Return type:

bool

load_state_dict(state_dict, **kwargs)

Updates registered buffers and loads parameters.

Parameters:

state_dict (dict) – dict containing parameters and persistent buffers

class sc2bench.models.layer.MSHPBasedResNetBottleneck(num_input_channels=3, num_latent_channels=16, num_bottleneck_channels=24, num_target_channels=256, g_a_channel_sizes=None, g_s_channel_sizes=None)

Mean-Scale Hyperprior(MSHP)-based bottleneck for ResNet and variants.

• David Minnen, Johannes Ballé, George Toderici: “Joint Autoregressive and Hierarchical Priors for Learned Image Compression” @ NeurIPS 2018 (2018)

• Yoshitomo Matsubara, Ruihan Yang, Marco Levorato, Stephan Mandt: “SC2 Benchmark: Supervised Compression for Split Computing” @ TMLR (2023)

Parameters:

• num_input_channels (int) – number of input channels

• num_latent_channels (int) – number of latent channels

• num_bottleneck_channels (int) – number of bottleneck channels

• num_target_channels (int) – number of output channels for decoder’s output

• g_a_channel_sizes (list[int] or None) – list of 4 numbers of channels for parametric transform \(g_a\)

• g_s_channel_sizes (list[int] or None) – list of 4 numbers of channels for parametric transform \(g_s\)

encode(x, **kwargs)

Encodes input data.

Parameters:

x (torch.Tensor) – input data

Returns:

entropy-coded compressed data (‘strings’ as key) and shape of the input data (‘shape’ as key)

Return type:

dict

decode(strings, shape)

Decodes encoded data.

Parameters:

• strings (list[str]) – entropy-coded compressed data

• shape (list[int]) – shape of the input data

Returns:

decompressed data

Return type:

torch.Tensor

sc2bench.models.layer.get_layer(cls_or_func_name, **kwargs)

Gets a layer module.

Parameters:

• cls_or_func_name (str) – layer class or function name

• kwargs (dict) – kwargs for the layer class or function to build a layer

Returns:

layer module

Return type:

nn.Module or None

sc2bench.models.registry.register_compressai_model(cls_or_func)

Registers a compression model class or a function to build a compression model in compressai.

Parameters:

cls_or_func (class or Callable) – compression model or function to build a compression model to be registered

Returns:

registered compression model class or function

Return type:

class or Callable

sc2bench.models.registry.register_compression_model_class(cls)

Registers a compression model class.

Parameters:

cls (class) – compression model to be registered

Returns:

registered compression model class

Return type:

class

sc2bench.models.registry.register_compression_model_func(func)

Registers a function to build a compression model.

Parameters:

func (Callable) – function to build a compression model to be registered

Returns:

registered function

Return type:

Callable

sc2bench.models.registry.get_compressai_model(compression_model_name, ckpt_file_path=None, updates=False, **compression_model_kwargs)

Gets a model in compressai.

Parameters:

• compression_model_name (str) – compressai model name

• ckpt_file_path (str or None) – checkpoint file path

• updates (bool) – if True, updates the parameters for entropy coding

• compression_model_kwargs (dict) – kwargs for the model class or function to build the model

Returns:

compressai model

Return type:

nn.Module

sc2bench.models.registry.get_compression_model(compression_model_config, device)

Gets a compression model.

Parameters:

• compression_model_config (dict) – compression model configuration

• device (str or torch.device) – torch device

Returns:

compression model

Return type:

nn.Module

sc2bench.models.registry.load_classification_model(model_config, device, distributed, strict=True)

Loads an image classification model.

Parameters:

• model_config (dict) – image classification model configuration

• device (str or torch.device) – torch device

• distributed (bool) – whether to use the model in distributed training mode

• strict (bool) – whether to strictly enforce that the keys in state_dict match the keys returned by the model’s state_dict() function

Returns:

image classification model

Return type:

nn.Module

sc2bench.models.wrapper.register_wrapper_class(cls)

Registers a model wrapper class.

Parameters:

cls (class) – model wrapper to be registered

Returns:

registered model wrapper class

Return type:

class

class sc2bench.models.wrapper.CodecInputCompressionClassifier(classification_model, device, codec_params=None, post_transform_params=None, analysis_config=None, **kwargs)

A wrapper module for codec input compression model followed by a classification model.

Parameters:

• classification_model (nn.Module) – image classification model

• device (torch.device or str) – torch device

• codec_params (dict or None) – transform sequence configuration for codec

• post_transform_params (dict or None) – post-transform parameters

• analysis_config (dict or None) – analysis configuration

class sc2bench.models.wrapper.NeuralInputCompressionClassifier(classification_model, pre_transform_params=None, compression_model=None, uses_cpu4compression_model=False, post_transform_params=None, analysis_config=None, **kwargs)

A wrapper module for neural input compression model followed by a classification model.

Parameters:

• classification_model (nn.Module) – image classification model

• pre_transform_params (dict or None) – pre-transform parameters

• compression_model (nn.Module or None) – compression model

• uses_cpu4compression_model (bool) – whether to use CPU instead of GPU for comoression_model

• post_transform_params (dict or None) – post-transform parameters

• analysis_config (dict or None) – analysis configuration

use_cpu4compression()

Changes the device of the compression model to CPU.

class sc2bench.models.wrapper.CodecFeatureCompressionClassifier(classification_model, device, encoder_config=None, codec_params=None, decoder_config=None, classifier_config=None, post_transform_params=None, analysis_config=None, **kwargs)

A wrapper module for codec feature compression model injected to a classification model.

Parameters:

• classification_model (nn.Module) – image classification model

• device (torch.device or str) – torch device

• encoder_config (dict or None) – configuration to design an encoder using modules in classification_model

• codec_params (dict or None) – transform sequence configuration for codec

• decoder_config (dict or None) – configuration to design a decoder using modules in classification_model

• classifier_config (dict or None) – configuration to design a classifier using modules in classification_model

• post_transform_params (dict or None) – post-transform parameters

• analysis_config (dict or None) – analysis configuration

class sc2bench.models.wrapper.EntropicClassifier(classification_model, encoder_config, compression_model_params, decoder_config, classifier_config, analysis_config=None, **kwargs)

A wrapper module for entropic compression model injected to a classification model.

Parameters:

• classification_model (nn.Module) – image classification model

• device (torch.device or str) – torch device

• encoder_config (dict) – configuration to design an encoder using modules in classification_model

• compression_model_params (dict) – kwargs for EntropyBottleneckLayer in compressai

• decoder_config (dict) – configuration to design a decoder using modules in classification_model

• classifier_config (dict) – configuration to design a classifier using modules in classification_model

• analysis_config (dict or None) – analysis configuration

update()

Updates compression-specific parameters like CompressAI models do.

This should be overridden by all subclasses.

load_state_dict(state_dict, **kwargs)

Loads parameters for all the sub-modules except entropy_bottleneck and then entropy_bottleneck.

Parameters:

state_dict (dict) – dict containing parameters and persistent buffers

get_aux_module(**kwargs)

Returns an auxiliary module to compute auxiliary loss if necessary like CompressAI models do.

This should be overridden by all subclasses.

class sc2bench.models.wrapper.SplitClassifier(classification_model, encoder_config, decoder_config, classifier_config, compressor_transform_params=None, decompressor_transform_params=None, analysis_config=None, **kwargs)

A wrapper module for naively splitting a classification model.

Parameters:

• classification_model (nn.Module) – image classification model

• encoder_config (dict or None) – configuration to design an encoder using modules in classification_model

• decoder_config (dict or None) – configuration to design a decoder using modules in classification_model

• classifier_config (dict or None) – configuration to design a classifier using modules in classification_model

• compressor_transform_params (dict or None) – transform parameters for compressor

• decompressor_transform_params (dict or None) – transform parameters for decompressor

• analysis_config (dict or None) – analysis configuration

update()

Updates compression-specific parameters like CompressAI models do.

This should be overridden by all subclasses.

get_aux_module(**kwargs)

Returns an auxiliary module to compute auxiliary loss if necessary like CompressAI models do.

This should be overridden by all subclasses.

sc2bench.models.wrapper.wrap_model(wrapper_model_name, model, compression_model, **kwargs)

Wraps a model and a compression model with a wrapper module.

Parameters:

• wrapper_model_name (str) – wrapper model name

• model (nn.Module) – model

• compression_model (nn.Module) – compression model

• kwargs (dict) – kwargs for the wrapper class or function to build the wrapper module

Returns:

wrapped model

Return type:

nn.Module

sc2bench.models.wrapper.get_wrapped_classification_model(wrapper_model_config, device, distributed)

Gets a wrapped image classification model.

Parameters:

• wrapper_model_config (dict) – wrapper model configuration

• device (torch.device) – torch device

• distributed (bool) – whether to use the model in distributed training mode

Returns:

wrapped image classification model

Return type:

nn.Module

sc2bench.models.detection

class sc2bench.models.detection.base.UpdatableDetectionModel(analyzer_configs=None)

An abstract class for updatable object detection model.

Parameters:

analyzer_configs (list[dict]) – list of analysis configurations

update(**kwargs)

Updates compression-specific parameters like CompressAI models do.

This should be overridden by all subclasses.

get_aux_module(**kwargs)

Returns an auxiliary module to compute auxiliary loss if necessary like CompressAI models do.

This should be overridden by all subclasses.

class sc2bench.models.detection.base.UpdatableBackboneWithFPN(backbone: Module, return_layer_dict: Dict[str, str], in_channels_list: List[int], out_channels: int, extra_blocks: Optional[ExtraFPNBlock] = None, analyzer_configs: Optional[List[Dict]] = None, analyzes_after_compress: bool = False, analyzable_layer_key: Optional[str] = None)

An updatable backbone model with feature pyramid network (FPN).

Parameters:

• backbone (nn.Module) – backbone model (usually a classification model)

• return_layer_dict (dict) – mapping from name of module to return its output to a specified key

• in_channels_list (list[int]) – number of channels for each feature map that is passed to the module for FPN

• out_channels (int) – number of channels of the FPN representation

• extra_blocks (ExtraFPNBlock or None) – if provided, extra operations will be performed. It is expected to take the fpn features, the original features and the names of the original features as input, and returns a new list of feature maps and their corresponding names

• analyzer_configs (list[dict]) – list of analysis configurations

• analyzes_after_compress (bool) – run analysis with analyzer_configs if True

• analyzable_layer_key (str or None) – key of analyzable layer

check_if_updatable()

Checks if this module is updatable with respect to CompressAI modules.

Returns:

True if the model is updatable, False otherwise

Return type:

bool

update()

Updates compression-specific parameters like CompressAI models do. Needs to be called once after training to be able to later perform the evaluation with an actual entropy coder.

get_aux_module()

Returns an auxiliary module to compute auxiliary loss if necessary like CompressAI models do.

Returns:

auxiliary module

Return type:

nn.Module

sc2bench.models.detection.base.check_if_updatable_detection_model(model)

Checks if the given object detection model is updatable.

Parameters:

model (nn.Module) – object detection model

Returns:

True if the model is updatable, False otherwise

Return type:

bool

class sc2bench.models.detection.rcnn.BaseRCNN(rcnn_model, analysis_config=None)

A base, updatable R-CNN model.

Parameters:

• rcnn_model (nn.Module) – backbone model (usually a classification model)

• analysis_config (dict or None) – analysis configuration

update(**kwargs)

Updates compression-specific parameters like CompressAI models do. Needs to be called once after training to be able to later perform the evaluation with an actual entropy coder.

get_aux_module(**kwargs)

Returns an auxiliary module to compute auxiliary loss if necessary like CompressAI models do.

Returns:

auxiliary module

Return type:

nn.Module

activate_analysis()

Activates the analysis mode.

Should be called after training model.

deactivate_analysis()

Deactivates the analysis mode.

analyze(compressed_obj)

Analyzes a given compressed object (e.g., file size of the compressed object).

Parameters:

compressed_obj (Any) – compressed object to be analyzed

summarize()

Summarizes the results that the configured analyzers store.

clear_analysis()

Clears the results that the configured analyzers store.

sc2bench.models.detection.rcnn.create_faster_rcnn_fpn(backbone, extra_blocks=None, return_layer_dict=None, in_channels_list=None, in_channels_stage2=None, out_channels=256, returned_layers=None, num_classes=91, analysis_config=None, analyzable_layer_key=None, **kwargs)

Builds Faster R-CNN model using a given updatable backbone model.

Parameters:

• backbone (nn.Module) – backbone model (usually a classification model)

• extra_blocks (ExtraFPNBlock or None) – if provided, extra operations will be performed. It is expected to take the fpn features, the original features and the names of the original features as input, and returns a new list of feature maps and their corresponding names

• return_layer_dict (dict) – mapping from name of module to return its output to a specified key

• in_channels_list (list[int] or None) – number of channels for each feature map that is passed to the module for FPN

• in_channels_stage2 (int or None) – base number of channels used to define in_channels_list if in_channels_list is None

• out_channels (int) – number of channels of the FPN representation

• returned_layers (list[int] or None) – list of layer numbers to define return_layer_dict if return_layer_dict is None

• num_classes (int) – number of output classes of the model (including the background)

• analysis_config (dict or None) – analysis configuration

• analyzable_layer_key (str or None) – key of analyzable layer

Returns:

Faster R-CNN model with backbone model with FPN

Return type:

torchvision.models.detection.faster_rcnn.FasterRCNN

sc2bench.models.detection.rcnn.faster_rcnn_model(backbone_config, pretrained=True, pretrained_backbone_name=None, progress=True, backbone_fpn_kwargs=None, num_classes=91, analysis_config=None, start_ckpt_file_path=None, **kwargs)

Builds Faster R-CNN model.

Parameters:

• backbone_config (dict) – backbone configuration

• pretrained (bool) – if True, returns a model pre-trained on COCO train2017 (torchvision)

• pretrained_backbone_name (str) – pretrained backbone name such as ‘resnet50’, ‘mobilenet_v3_large_320’, and ‘mobilenet_v3_large’

• progress (bool) – if True, displays a progress bar of the download to stderr

• backbone_fpn_kwargs (dict) – keyword arguments for create_faster_rcnn_fpn

• num_classes (int) – number of output classes of the model (including the background)

• analysis_config (dict or None) – analysis configuration

• start_ckpt_file_path (str or None) – checkpoint file path to be loaded for the built Faster R-CNN model

Returns:

Faster R-CNN model with splittable backbone model and FPN

Return type:

BaseRCNN

sc2bench.models.detection.registry.register_detection_model_class(cls)

Registers an object detection model class.

Parameters:

cls (class) – object detection model class to be registered

Returns:

registered object detection model class

Return type:

class

sc2bench.models.detection.registry.register_detection_model_func(func)

Registers a function to build an object detection model.

Parameters:

func (Callable) – function to build an object detection model to be registered

Returns:

registered function

Return type:

Callable

sc2bench.models.detection.registry.get_detection_model(cls_or_func_name, **kwargs)

Gets an object detection model.

Parameters:

cls_or_func_name (str) – model class or function name

Returns:

object detection model

Return type:

nn.Module or None

sc2bench.models.detection.registry.load_detection_model(model_config, device, strict=True)

Loads an object detection model.

Parameters:

• model_config (dict) – model configuration

• device (torch.device or str) – device

• strict (bool) – whether to strictly enforce that the keys in state_dict match the keys returned by this module’s state_dict() function.

Returns:

object detection model

Return type:

nn.Module

class sc2bench.models.detection.transform.RCNNTransformWithCompression(transform, device, codec_params, analyzer_configs, analyzes_after_compress=False, compression_model=None, uses_cpu4compression_model=False, pre_transform_params=None, post_transform_params=None, adaptive_pad_kwargs=None)

An R-CNN Transform with codec-based or model-based compression.

Parameters:

• transform (nn.Module) – performs the data transformation from the inputs to feed into the model

• device (torch.device or str) – torch device

• codec_params (dict) – codec parameters

• analyzer_configs (list[dict]) – list of analysis configurations

• analyzes_after_compress (bool) – run analysis with analyzer_configs if True

• compression_model (nn.Module or None) – compression model

• uses_cpu4compression_model (bool) – whether to use CPU instead of GPU for comoression_model

• pre_transform_params (dict or None) – pre-transform parameters

• post_transform_params (dict or None) – post-transform parameters

• adaptive_pad_kwargs (dict or None) – keyword arguments for AdaptivePad

compress_by_codec(org_img)

Convert a tensor to an image and compress-decompress it by codec.

Parameters:

org_img (torch.Tensor) – image tensor

Returns:

compressed-and-decompressed image tensor

Return type:

torch.Tensor

compress_by_model(org_img)

Convert a tensor to an image and compress-decompress it by model.

Parameters:

org_img (torch.Tensor) – image tensor

Returns:

compressed-and-decompressed image tensor

Return type:

torch.Tensor

compress(org_img)

Apply pre_transform to an image tensor, compress and decompress it, and apply post_transform to the compressed-decompressed image tensor.

Parameters:

org_img (torch.Tensor) – image tensor

Returns:

compressed-and-decompressed image tensor

Return type:

torch.Tensor

class sc2bench.models.detection.wrapper.InputCompressionDetectionModel(detection_model, device, codec_params=None, compression_model=None, uses_cpu4compression_model=False, pre_transform_params=None, post_transform_params=None, analysis_config=None, adaptive_pad_kwargs=None, **kwargs)

A wrapper module for input compression model followed by a detection model.

Parameters:

• detection_model (nn.Module) – object detection model

• device (torch.device or str) – torch device

• codec_params (dict or None) – transform sequence configuration for codec

• compression_model (nn.Module or None) – compression model

• uses_cpu4compression_model (bool) – whether to use CPU instead of GPU for comoression_model

• pre_transform_params (dict or None) – pre-transform parameters

• post_transform_params (dict or None) – post-transform parameters

• analysis_config (dict or None) – analysis configuration

• adaptive_pad_kwargs (dict or None) – keyword arguments for AdaptivePad

use_cpu4compression()

Changes the device of the compression model to CPU.

activate_analysis()

Makes internal analyzers ready to run.

deactivate_analysis()

Turns internal analyzers off.

analyze(compressed_obj)

Analyzes a compressed object using internal analyzers.

Parameters:

compressed_obj (Any) – compressed object

summarize()

Shows each of internal analyzers’ summary of results.

clear_analysis()

Clears each of internal analyzers’ results.

sc2bench.models.detection.wrapper.get_wrapped_detection_model(wrapper_model_config, device)

Gets a wrapped object detection model.

Parameters:

• wrapper_model_config (dict) – wrapper model configuration

• device (torch.device) – torch device

Returns:

wrapped object detection model

Return type:

nn.Module

sc2bench.models.segmentation

class sc2bench.models.segmentation.base.UpdatableSegmentationModel(analyzer_configs=None)

An abstract class for updatable semantic segmentation model.

Parameters:

analyzer_configs (list[dict]) – list of analysis configurations

update(**kwargs)

Updates compression-specific parameters like CompressAI models do.

This should be overridden by all subclasses.

get_aux_module(**kwargs)

Returns an auxiliary module to compute auxiliary loss if necessary like CompressAI models do.

This should be overridden by all subclasses.

class sc2bench.models.segmentation.base.BaseSegmentationModel(backbone, classifier, aux_classifier=None, analysis_config=None)

update(**kwargs)

Updates compression-specific parameters like CompressAI models do. Needs to be called once after training to be able to later perform the evaluation with an actual entropy coder.

get_aux_module(**kwargs)

Returns an auxiliary module to compute auxiliary loss if necessary like CompressAI models do.

Returns:

auxiliary module

Return type:

nn.Module

activate_analysis()

Activates the analysis mode.

Should be called after training model.

deactivate_analysis()

Deactivates the analysis mode.

analyze(compressed_obj)

Analyzes a given compressed object (e.g., file size of the compressed object).

Parameters:

compressed_obj (Any) – compressed object to be analyzed

summarize()

Summarizes the results that the configured analyzers store.

clear_analysis()

Clears the results that the configured analyzers store.

sc2bench.models.segmentation.base.check_if_updatable_segmentation_model(model)

Checks if the given semantic segmentation model is updatable.

Parameters:

model (nn.Module) – semantic segmentation model

Returns:

True if the model is updatable, False otherwise

Return type:

bool

sc2bench.models.segmentation.deeplabv3.create_deeplabv3(backbone, num_input_channels=2048, uses_aux=False, num_aux_channels=1024, num_classes=21)

Builds DeepLabv3 model using a given updatable backbone model.

Parameters:

• backbone (nn.Module) – backbone model (usually a classification model)

• num_input_channels (int) – number of input channels for classification head

• uses_aux (bool) – If True, add an auxiliary branch

• num_aux_channels (int) – number of input channels for auxiliary classification head

• num_classes (int) – number of output classes of the model (including the background)

Returns:

DeepLabv3 model

Return type:

BaseSegmentationModel

sc2bench.models.segmentation.deeplabv3.deeplabv3_model(backbone_config, pretrained=True, pretrained_backbone_name=None, progress=True, num_input_channels=2048, uses_aux=False, num_aux_channels=1024, return_layer_dict=None, num_classes=21, analysis_config=None, analyzable_layer_key=None, start_ckpt_file_path=None, **kwargs)

Builds DeepLabv3 model using a given updatable backbone model.

Parameters:

• backbone_config (dict) – backbone configuration

• pretrained (bool) – if True, returns a model pre-trained on COCO train2017 (torchvision)

• pretrained_backbone_name (str) – pretrained backbone name such as ‘resnet50’, ‘resnet101’, and ‘mobilenet_v3_large’

• progress (bool) – if True, displays a progress bar of the download to stderr

• num_input_channels (int) – number of input channels for classification head

• uses_aux (bool) – If True, add an auxiliary branch

• num_aux_channels (int) – number of input channels for auxiliary classification head

• return_layer_dict (dict) – mapping from name of module to return its output to a specified key

• num_classes (int) – number of output classes of the model (including the background)

• analysis_config (dict or None) – analysis configuration

• analyzable_layer_key (str or None) – key of analyzable layer

• start_ckpt_file_path (str or None) – checkpoint file path to be loaded for the built DeepLabv3 model

Returns:

DeepLabv3 model with splittable backbone model

Return type:

BaseSegmentationModel

sc2bench.models.segmentation.registry.register_segmentation_model_class(cls)

Registers a semantic segmentation model class.

Parameters:

cls (class) – semantic segmentation model to be registered

Returns:

registered semantic segmentation model class

Return type:

class

sc2bench.models.segmentation.registry.register_segmentation_model_func(func)

Registers a function to build a semantic segmentation model.

Parameters:

func (Callable) – function to build a semantic segmentation model to be registered

Returns:

registered function

Return type:

Callable

sc2bench.models.segmentation.registry.get_segmentation_model(cls_or_func_name, **kwargs)

Gets a semantic segmentation model.

Parameters:

cls_or_func_name (str) – model class or function name

Returns:

semantic segmentation model

Return type:

nn.Module or None

sc2bench.models.segmentation.registry.load_segmentation_model(model_config, device, strict=True)

Loads a semantic segmentation model.

Parameters:

• model_config (dict) – model configuration

• device (torch.device or str) – device

• strict (bool) – whether to strictly enforce that the keys in state_dict match the keys returned by this module’s state_dict() function.

Returns:

semantic segmentation model

Return type:

nn.Module

class sc2bench.models.segmentation.wrapper.CodecInputCompressionSegmentationModel(segmentation_model, device, codec_params=None, post_transform_params=None, analysis_config=None, **kwargs)

A wrapper module for codec input compression model followed by a segmentation model.

Parameters:

• segmentation_model (nn.Module) – semantic segmentation model

• device (torch.device or str) – torch device

• codec_params (dict or None) – transform sequence configuration for codec

• post_transform_params (dict or None) – post-transform parameters

• analysis_config (dict or None) – analysis configuration

class sc2bench.models.segmentation.wrapper.NeuralInputCompressionSegmentationModel(segmentation_model, pre_transform_params=None, compression_model=None, uses_cpu4compression_model=False, post_transform_params=None, analysis_config=None, **kwargs)

A wrapper module for neural input compression model followed by a segmentation model.

Parameters:

• segmentation_model (nn.Module) – semantic segmentation model

• pre_transform_params (dict or None) – pre-transform parameters

• compression_model (nn.Module or None) – compression model

• uses_cpu4compression_model (bool) – whether to use CPU instead of GPU for comoression_model

• post_transform_params (dict or None) – post-transform parameters

• analysis_config (dict or None) – analysis configuration

use_cpu4compression()

Changes the device of the compression model to CPU.

sc2bench.models.segmentation.wrapper.get_wrapped_segmentation_model(wrapper_model_config, device)

Gets a wrapped semantic segmentation model.

Parameters:

• wrapper_model_config (dict) – wrapper model configuration

• device (torch.device) – torch device

Returns:

wrapped semantic segmentation model

Return type:

nn.Module

sc2bench.transforms

sc2bench.transforms.codec.register_codec_transform_module(cls)

Registers a codec transform class.

Parameters:

cls (class) – codec transform class to be registered

Returns:

registered codec transform class

Return type:

class

class sc2bench.transforms.codec.WrappedRandomResizedCrop(interpolation=None, **kwargs)

RandomResizedCrop in torchvision wrapped to be defined by interpolation as a str object.

Parameters:

• interpolation (str or None) – desired interpolation mode (‘nearest’, ‘bicubic’, ‘bilinear’, ‘box’, ‘hamming’, ‘lanczos’)

• kwargs (dict) – kwargs for RandomResizedCrop in torchvision

class sc2bench.transforms.codec.WrappedResize(interpolation=None, **kwargs)

Resize in torchvision wrapped to be defined by interpolation as a str object.

Parameters:

• interpolation (str or None) – desired interpolation mode (‘nearest’, ‘bicubic’, ‘bilinear’, ‘box’, ‘hamming’, ‘lanczos’)

• kwargs (dict) – kwargs for Resize in torchvision

class sc2bench.transforms.codec.PILImageModule(returns_file_size=False, open_kwargs=None, **save_kwargs)

A generalized PIL module to compress (decompress) images e.g., as part of transform pipeline.

Parameters:

• returns_file_size (bool) – returns file size of compressed object in addition to PIL image if True

• open_kwargs (dict or None) – kwargs to be used as part of Image.open(img_buffer, **open_kwargs)

• save_kwargs (dict or None) – kwargs to be used as part of Image.save(img_buffer, **save_kwargs)

forward(pil_img, *args)

Saves PIL Image to BytesIO and reopens the image saved in the buffer.

Parameters:

pil_img (PIL.Image.Image) – image to be transformed.

Returns:

Affine transformed image or with its file size if returns_file_size=True

Return type:

PIL.Image.Image or (PIL.Image.Image, int)

class sc2bench.transforms.codec.PILTensorModule(returns_file_size=False, open_kwargs=None, **save_kwargs)

A generalized PIL module to compress (decompress) tensors e.g., as part of transform pipeline.

Parameters:

• returns_file_size (bool) – returns file size of compressed object in addition to PIL image if True

• open_kwargs (dict or None) – kwargs to be used as part of Image.open(img_buffer, **open_kwargs)

• save_kwargs (dict or None) – kwargs to be used as part of Image.save(img_buffer, **save_kwargs)

forward(x, *args)

Splits tensor’s channels into sub-tensors (3 or fewer channels each), normalizes each using its min and max values, saves the normalized sub-tensor to BytesIO, and reopens the sub-tensor saved in the buffer to reconstruct the input tensor.

Parameters:

x (torch.Tensor) – image tensor (C, H, W) to be transformed.

Returns:

Affine transformed image tensor or with its file size if returns_file_size=True

Return type:

torch.Tensor or (torch.Tensor, int)

class sc2bench.transforms.codec.BPGModule(encoder_path, decoder_path, color_mode='ycbcr', encoder='x265', subsampling_mode='444', bit_depth='8', quality=50, returns_file_size=False)

A BPG module to compress (decompress) images e.g., as part of transform pipeline.

Modified https://github.com/InterDigitalInc/CompressAI/blob/master/compressai/utils/bench/codecs.py

Fabrice Bellard: “BPG Image format”

Warning

You need to manually install BPG software beforehand and confirm the encoder and decoder paths. For Debian machines (e.g., Ubuntu), you can use this script.

Parameters:

• encoder_path (str) – file path of BPG encoder you manually installed

• decoder_path (str) – file path of BPG decoder you manually installed

• color_mode (str) – color mode (‘ycbcr’ or ‘rgb’)

• encoder (str) – encoder type (‘x265’ or ‘jctvc’)

• subsampling_mode (str or int) – subsampling mode (‘420’ or ‘444’)

• bit_depth (str or int) – bit depth (8 or 10)

• quality (int) – quality value in range [0, 51]

• returns_file_size (bool) – returns file size of compressed object in addition to PIL image if True

forward(pil_img)

Compresses and decompresses PIL Image using BPG software.

Parameters:

pil_img (PIL.Image.Image) – image to be transformed.

Returns:

Affine transformed image or with its file size if returns_file_size=True

Return type:

PIL.Image.Image or (PIL.Image.Image, int)

class sc2bench.transforms.codec.VTMModule(encoder_path, decoder_path, config_path, color_mode='ycbcr', quality=63, returns_file_size=False)

A VTM module to compress (decompress) images e.g., as part of transform pipeline.

Modified https://github.com/InterDigitalInc/CompressAI/blob/master/compressai/utils/bench/codecs.py

The Joint Video Exploration Team: “VTM reference software for VVC”

Warning

You need to manually install VTM software beforehand and confirm the encoder and decoder paths. For Debian machines (e.g., Ubuntu), you can use this script.

Parameters:

• encoder_path (str) – file path of VTM encoder you manually installed

• decoder_path (str) – file path of VTM decoder you manually installed

• config_path (str) – VTM configuration file path

• color_mode (str) – color mode (‘ycbcr’ or ‘rgb’)

• quality (int) – quality value in range [0, 63]

• returns_file_size (bool) – returns file size of compressed object in addition to PIL image if True

forward(pil_img)

Compresses and decompresses PIL Image using VTM software.

Parameters:

pil_img (PIL.Image.Image) – image to be transformed.

Returns:

Affine transformed image or with its file size if returns_file_size=True

Return type:

PIL.Image.Image or (PIL.Image.Image, int)

sc2bench.transforms.collator.cat_list(images, fill_value=0)

Concatenates a list of images with the max size for each of heights and widths and fills empty spaces with a specified value.

Parameters:

• images (torch.Tensor) – batch tensor

• fill_value (int) – value to be filled

Returns:

backbone model

Return type:

torch.Tensor

sc2bench.transforms.collator.pascal_seg_collate_fn(batch)

Collates input data for PASCAL VOC 2012 segmentation.

Parameters:

batch (list or tuple) – list/tuple of triplets (image, target, supp_dict), where supp_dict can be an empty dict

Returns:

collated images, targets, and supplementary dicts

Return type:

(torch.Tensor, tensor.Tensor, list[dict])

sc2bench.transforms.collator.pascal_seg_eval_collate_fn(batch)

Collates input data for PASCAL VOC 2012 segmentation in evaluation

Parameters:

batch (list or tuple) – list/tuple of tuples (image, target)

Returns:

collated images and targets

Return type:

(torch.Tensor, tensor.Tensor)

sc2bench.transforms.misc.register_misc_transform_module(cls)

Registers a miscellaneous transform class.

Parameters:

cls (class) – miscellaneous transform class to be registered

Returns:

registered miscellaneous transform class

Return type:

class

sc2bench.transforms.misc.default_collate_w_pil(batch)

Puts each data field into a tensor or PIL Image with outer dimension batch size.

Parameters:

batch – single batch to be collated

Returns:

collated batch

class sc2bench.transforms.misc.ClearTargetTransform

A transform module that replaces target with an empty list.

forward(sample, *args)

Replaces target data field with an empty list.

Parameters:

sample (PIL.Image.Image or torch.Tensor) – image or image tensor

Returns:

sample and an empty list

Return type:

(PIL.Image.Image or torch.Tensor, list)

class sc2bench.transforms.misc.AdaptivePad(fill=0, padding_position='hw', padding_mode='constant', factor=128, returns_org_patch_size=False)

A transform module that adaptively determines the size of padded sample.

Parameters:

• fill (int) – padded value

• padding_position (str) – ‘hw’ (default) to pad left and right for padding horizontal size // 2 and top and bottom for padding vertical size // 2; ‘right_bottom’ to pad bottom and right only

• padding_mode (str) – padding mode passed to pad module

• factor (int) – factor value for the padded input sample

• returns_org_patch_size (bool) – if True, returns the patch size of the original input

forward(x)

Adaptively determines the size of padded image or image tensor.

Parameters:

x (PIL.Image.Image or torch.Tensor) – image or image tensor

Returns:

padded image or image tensor, and the patch size of the input (height, width) if returns_org_patch_size=True

Return type:

PIL.Image.Image or torch.Tensor or (PIL.Image.Image or torch.Tensor, list[int, int])

class sc2bench.transforms.misc.CustomToTensor(converts_sample=True, converts_target=True)

A customized ToTensor module that can be applied to sample and target selectively.

Parameters:

• converts_sample (bool) – if True, applies to_tensor to sample

• converts_target (bool) – if True, applies torch.as_tensor to target

class sc2bench.transforms.misc.SimpleQuantizer(num_bits)

A module to quantize tensor with its half() function if num_bits=16 (FP16) or Jacob et al.’s method if num_bits=8 (INT8 + one FP32 scale parameter).

Benoit Jacob, Skirmantas Kligys, Bo Chen, Menglong Zhu, Matthew Tang, Andrew Howard, Hartwig Adam, Dmitry Kalenichenko: “Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference” @ CVPR 2018 (2018)

Parameters:

num_bits (int) – number of bits for quantization

forward(z)

Quantizes tensor.

Parameters:

z (torch.Tensor) – tensor

Returns:

quantized tensor

Return type:

torch.Tensor or torchdistill.common.tensor_util.QuantizedTensor

class sc2bench.transforms.misc.SimpleDequantizer(num_bits)

A module to dequantize quantized tensor in FP32. If num_bits=8, it uses Jacob et al.’s method.

Benoit Jacob, Skirmantas Kligys, Bo Chen, Menglong Zhu, Matthew Tang, Andrew Howard, Hartwig Adam, Dmitry Kalenichenko: “Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference” @ CVPR 2018 (2018)

Parameters:

num_bits (int) – number of bits used for quantization

forward(z)

Dequantizes quantized tensor.

Parameters:

z (torch.Tensor or torchdistill.common.tensor_util.QuantizedTensor) – quantized tensor

Returns:

dequantized tensor

Return type:

torch.Tensor

sc2bench.analysis

sc2bench.analysis.register_analysis_class(cls)

Registers an analyzer class.

Parameters:

cls (class) – analyzer class to be registered

Returns:

registered analyzer class

Return type:

class

class sc2bench.analysis.AnalyzableModule(analyzer_configs=None)

A base module to analyze and summarize the wrapped modules and intermediate representations.

Parameters:

analyzer_configs (list[dict] or None) – list of analysis configurations

forward(*args, **kwargs)

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.

activate_analysis()

Makes internal analyzers ready to run.

deactivate_analysis()

Turns internal analyzers off.

analyze(compressed_obj)

Analyzes a compressed object using internal analyzers.

Parameters:

compressed_obj (Any) – compressed object

summarize()

Shows each of internal analyzers’ summary of results.

clear_analysis()

Clears each of internal analyzers’ results.

class sc2bench.analysis.BaseAnalyzer

A base analyzer to analyze and summarize the wrapped modules and intermediate representations.

analyze(*args, **kwargs)

Analyzes a compressed object.

summarize()

Shows the summary of results.

This should be overridden by all subclasses.

clear()

Clears the results.

This should be overridden by all subclasses.

class sc2bench.analysis.FileSizeAnalyzer(unit='KB', **kwargs)

An analyzer to measure file size of compressed object in the designated unit.

Parameters:

unit (str) – unit of data size in bytes (‘B’, ‘KB’, ‘MB’)

analyze(compressed_obj)

Computes and appends binary object size of the compressed object.

Parameters:

compressed_obj (Any) – compressed object

summarize()

Computes and shows mean and std of the stored file sizes and the number of samples .

clear()

Clears the file size list.

class sc2bench.analysis.FileSizeAccumulator(unit='KB', **kwargs)

An accumulator to store pre-computed file size in the designated unit.

Parameters:

unit (str) – unit of data size in bytes (‘B’, ‘KB’, ‘MB’)

analyze(file_size)

Appends a file size.

Parameters:

file_size (int or float) – pre-computed file size

sc2bench.analysis.get_analyzer(cls_name, **kwargs)

Gets an analyzer module.

Parameters:

• cls_name (str) – analyzer class name

• kwargs (dict) – kwargs for the analyzer class

Returns:

analyzer module

Return type:

BaseAnalyzer or None

sc2bench.analysis.check_if_analyzable(module)

Checks if a module is an instance of AnalyzableModule.

Parameters:

module (Any) – module

Returns:

True if the module is an instance of AnalyzableModule. False otherwise

Return type:

bool

sc2bench.analysis.analyze_model_size(model, encoder_paths=None, additional_rest_paths=None, ignores_dtype_error=True)

Approximates numbers of bits used for parameters of the whole model, encoder, and the rest of the model.

Parameters:

• model (nn.Module) – model

• encoder_paths (list[str] or None) – list of module paths for the model to be considered as part of encoder’s parameters

• additional_rest_paths (list[str] or None) – list of additional rest module paths whose parameters should be shared with encoder e.g., module path of entropy bottleneck in the model if applied

• ignores_dtype_error (bool) – if False, raise an error when any unexpected dtypes are found

Returns:

model size (sum of param x num_bits) with three keys: model (whole model), encoder, and the rest

Return type:

dict

sc2bench.loss

sc2bench.loss.extract_org_loss_dict(org_criterion, student_outputs, teacher_outputs, targets, uses_teacher_output, **kwargs)

Extracts loss(es) from student_outputs inside TrainingBox or DistillationBox in torchdistill.

Parameters:

• org_criterion (nn.Module) – not used

• student_outputs (dict or Any) – student models’ output

• teacher_outputs (Any) – not used

• targets (Any) – not used

• uses_teacher_output (bool) – not used

Returns:

original loss dict

Return type:

class

sc2bench.loss.extract_org_segment_loss(org_criterion, student_outputs, teacher_outputs, targets, uses_teacher_output, **kwargs)

Computes loss(es) using the original loss module inside TrainingBox or DistillationBox in torchdistill for semantic segmentation models in torchvision.

Parameters:

• org_criterion (nn.Module) – original loss module

• student_outputs (dict or Any) – student models’ output

• teacher_outputs (Any) – not used

• targets (Any) – targets

• uses_teacher_output (bool) – not used

Returns:

original loss dict

Return type:

class

class sc2bench.loss.BppLoss(entropy_module_path, reduction='mean')

Bit-per-pixel (or rate) loss.

Parameters:

• entropy_module_path (str) – entropy module path to extract its output from io_dict

• reduction (str or None) – reduction type (‘sum’, ‘batchmean’, or ‘mean’)

forward(student_io_dict, *args, **kwargs)

Computes a rate loss.

Parameters:

student_io_dict (dict) – io_dict of model to be trained