sc2bench.models

---

sc2bench.models.layer

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=7, compressor_transform=None, decompressor_transform=None)

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

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

Yoshitomo Matsubara, Davide Callegaro, Sabur Baidya, Marco Levorato, Sameer Singh: “Head Network Distillation: Splitting Distilled Deep Neural Networks for Resource-constrained Edge Computing Systems” @ IEEE Access (2020)

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)

• compressor_transform (nn.Module or None) – compressor transform

• decompressor_transform (nn.Module or None) – decompressor transform

Returns:

bottleneck layer consisting of encoder and decoder

Return type:

SimpleBottleneck

sc2bench.models.layer.larger_densenet_bottleneck(bottleneck_channel=12, bottleneck_idx=8, compressor_transform=None, decompressor_transform=None)

Builds a bottleneck layer DenseNet-based encoder and decoder (23 layers in total).

Compatible with DenseNet-169 and -201.

Yoshitomo Matsubara, Davide Callegaro, Sabur Baidya, Marco Levorato, Sameer Singh: “Head Network Distillation: Splitting Distilled Deep Neural Networks for Resource-constrained Edge Computing Systems” @ IEEE Access (2020)

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)

• compressor_transform (nn.Module or None) – compressor transform

• decompressor_transform (nn.Module or None) – decompressor transform

Returns:

bottleneck layer consisting of encoder and decoder

Return type:

SimpleBottleneck

sc2bench.models.layer.inception_v3_bottleneck(bottleneck_channel=12, bottleneck_idx=7, compressor_transform=None, decompressor_transform=None)

Builds a bottleneck layer InceptionV3-based encoder and decoder (17 layers in total).

Yoshitomo Matsubara, Davide Callegaro, Sabur Baidya, Marco Levorato, Sameer Singh: “Head Network Distillation: Splitting Distilled Deep Neural Networks for Resource-constrained Edge Computing Systems” @ IEEE Access (2020)

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)

• compressor_transform (nn.Module or None) – compressor transform

• decompressor_transform (nn.Module or None) – decompressor transform

Returns:

bottleneck layer consisting of encoder and decoder

Return type:

SimpleBottleneck

sc2bench.models.layer.smaller_resnet_layer1_bottleneck(bottleneck_channel=12, bottleneck_idx=8, compressor_transform=None, decompressor_transform=None, **kwargs)

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

Compatible with the first layers of ResNet-18 and -34 until layer1.

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)

• compressor_transform (nn.Module or None) – compressor transform

• decompressor_transform (nn.Module or None) – decompressor transform

Returns:

bottleneck layer consisting of encoder and decoder

Return type:

SimpleBottleneck

sc2bench.models.layer.larger_resnet_layer1_bottleneck(bottleneck_channel=12, bottleneck_idx=8, compressor_transform=None, decompressor_transform=None, **kwargs)

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

Compatible with the first layers of ResNet-50, -101, and -152 until layer1.

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)

• compressor_transform (nn.Module or None) – compressor transform

• decompressor_transform (nn.Module or None) – decompressor transform

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

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 or None) – 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

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_encoder_decoder=None, post_transform=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_encoder_decoder (nn.Module or None) – transform sequence configuration for codec

• post_transform (nn.Module or None) – post-transform

• analysis_config (dict or None) – analysis configuration

class sc2bench.models.wrapper.NeuralInputCompressionClassifier(classification_model, pre_transform=None, compression_model=None, uses_cpu4compression_model=False, post_transform=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 (nn.Module or None) – pre-transform

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

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

• post_transform (nn.Module or None) – post-transform

• 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_encoder_decoder=None, decoder_config=None, classifier_config=None, post_transform=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_encoder_decoder (nn.Module 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 (nn.Module or None) – post-transform

• analysis_config (dict or None) – analysis configuration

class sc2bench.models.wrapper.EntropicClassifier(classification_model, encoder_config, compression_model_kwargs, 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_kwargs (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=None, decompressor_transform=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 (nn.Module or None) – compressor transform

• decompressor_transform (nn.Module or None) – decompressor transform

• 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.backbone

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=None, analysis_config=None, short_module_names=None)

ResNet/ResNeSt-based splittable image classification model optionally 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 (nn.Module or None) – pre-transform

• analysis_config (dict or None) – analysis configuration

• short_module_names (list) – child module names of “ResNet” to use

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.SplittableDenseNet(bottleneck_layer, short_feature_names, densenet_model, skips_avgpool=True, skips_classifier=True, pre_transform=None, analysis_config=None)

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

• Gao Huang, Zhuang Liu, Laurens van der Maaten, Kilian Q. Weinberger: “Densely Connected Convolutional Networks” @ CVPR 2017 (2017)

• Yoshitomo Matsubara, Davide Callegaro, Sabur Baidya, Marco Levorato, Sameer Singh: “Head Network Distillation: Splitting Distilled Deep Neural Networks for Resource-constrained Edge Computing Systems” @ IEEE Access (2020)

Parameters:

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

• short_feature_names (list) – child module names of “features” to use

• densenet_model (nn.Module) – DenseNet model to be used as a base model

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

• skips_classifier (bool) – if True, skips classifier after average pooling

• pre_transform (nn.Module or None) – pre-transform

• 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.SplittableInceptionV3(bottleneck_layer, short_module_names, inception_v3_model, skips_avgpool=True, skips_dropout=True, skips_fc=True, pre_transform=None, analysis_config=None)

Inception V3-based splittable image classification model optionally containing neural encoder, entropy bottleneck, and decoder.

• Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, Zbigniew Wojna: “Rethinking the Inception Architecture for Computer Vision” @ CVPR 2016 (2016)

• Yoshitomo Matsubara, Davide Callegaro, Sabur Baidya, Marco Levorato, Sameer Singh: “Head Network Distillation: Splitting Distilled Deep Neural Networks for Resource-constrained Edge Computing Systems” @ IEEE Access (2020)

Parameters:

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

• short_module_names (list) – child module names of “Inception3” to use

• inception_v3_model (nn.Module) – Inception V3 model to be used as a base model

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

• skips_dropout (bool) – if True, skips dropout after avgpool

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

• pre_transform (nn.Module or None) – pre-transform

• 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=None, analysis_config=None)

RegNet-based splittable image classification model optionally 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 (nn.Module or None) – pre-transform

• 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=None, analysis_config=None)

Hybrid ViT-based splittable image classification model optionally 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 (nn.Module or None) – pre-transform

• 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=None, analysis_config=None, org_model_ckpt_file_path_or_url=None, org_ckpt_strict=True, short_module_names=None, **resnet_kwargs)

Builds ResNet-based splittable image classification model optionally 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 (nn.Module or None) – pre-transform

• 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

• short_module_names (list) – child module names of “ResNet” to use

Returns:

splittable ResNet model

Return type:

SplittableResNet

sc2bench.models.backbone.splittable_densenet(bottleneck_config, densenet_name='densenet169', short_feature_names=None, skips_avgpool=True, skips_classifier=True, pre_transform=None, analysis_config=None, org_model_ckpt_file_path_or_url=None, org_ckpt_strict=True, **densenet_kwargs)

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

Parameters:

• bottleneck_config (dict) – bottleneck configuration

• densenet_name (str) – name of DenseNet function in torchvision

• short_feature_names (list) – child module names of “features” to use

• skips_avgpool (bool) – if True, skips adaptive_avgpool (average pooling) after features

• skips_classifier (bool) – if True, skips classifier after average pooling

• pre_transform (nn.Module or None) – pre-transform

• analysis_config (dict or None) – analysis configuration

• org_model_ckpt_file_path_or_url (str or None) – original DenseNet 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 DenseNet model’s state_dict() function

Returns:

splittable DenseNet model

Return type:

SplittableDenseNet

sc2bench.models.backbone.splittable_inception_v3(bottleneck_config, short_module_names=None, skips_avgpool=True, skips_dropout=True, skips_fc=True, pre_transform=None, analysis_config=None, org_model_ckpt_file_path_or_url=None, org_ckpt_strict=True, **inception_v3_kwargs)

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

Parameters:

• bottleneck_config (dict) – bottleneck configuration

• short_module_names (list) – child module names of “Inception3” to use

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

• skips_dropout (bool) – if True, skips dropout after avgpool

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

• pre_transform (nn.Module or None) – pre-transform

• analysis_config (dict or None) – analysis configuration

• org_model_ckpt_file_path_or_url (str or None) – original InceptionV3 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 InceptionV3 model’s state_dict() function

Returns:

splittable InceptionV3 model

Return type:

SplittableInceptionV3

sc2bench.models.backbone.splittable_resnest(bottleneck_config, resnest_name='resnest50d', inplanes=None, skips_avgpool=True, skips_fc=True, pre_transform=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 optionally 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 (nn.Module or None) – pre-transform

• 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=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 optionally 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 (nn.Module or None) – pre-transform

• 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=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 optionally 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 (nn.Module or None) – pre-transform

• 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.detection

---

sc2bench.models.detection.base

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: ExtraFPNBlock | None = None, analyzer_configs: List[Dict] | None = None, analyzes_after_compress: bool = False, analyzable_layer_key: str | None = 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

---

sc2bench.models.detection.rcnn

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

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

---

sc2bench.models.detection.transform

class sc2bench.models.detection.transform.RCNNTransformWithCompression(transform, device, codec_encoder_decoder, analyzer_configs, analyzes_after_compress=False, compression_model=None, uses_cpu4compression_model=False, pre_transform=None, post_transform=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_encoder_decoder (nn.Module) – transform sequence configuration for codec

• 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 compression_model

• pre_transform (nn.Module or None) – pre-transform

• post_transform (nn.Module or None) – post-transform

• 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

---

sc2bench.models.detection.wrapper

class sc2bench.models.detection.wrapper.InputCompressionDetectionModel(detection_model, device, codec_encoder_decoder=None, compression_model=None, uses_cpu4compression_model=False, pre_transform=None, post_transform=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_encoder_decoder (nn.Module 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 compression_model

• pre_transform (nn.Module or None) – pre-transform

• post_transform (nn.Module or None) – post-transform

• 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

---

sc2bench.models.segmentation.base

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

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

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

---

sc2bench.models.segmentation.wrapper

class sc2bench.models.segmentation.wrapper.CodecInputCompressionSegmentationModel(segmentation_model, device, codec_encoder_decoder=None, post_transform=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_encoder_decoder (nn.Module or None) – transform sequence configuration for codec

• post_transform (nn.Module or None) – post-transform

• analysis_config (dict or None) – analysis configuration

class sc2bench.models.segmentation.wrapper.NeuralInputCompressionSegmentationModel(segmentation_model, pre_transform=None, compression_model=None, uses_cpu4compression_model=False, post_transform=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 (nn.Module or None) – pre-transform

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

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

• post_transform (nn.Module or None) – post-transform

• 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