Architecture Templates

Architecture Templates: Pre-defined model architectures with customization.

Templates provide starting points based on well-known architectures, with full customization capabilities.

Example

>>> from torchvision_customizer.templates import resnet, vgg

>>> # Simple usage
>>> model = resnet(layers=50, num_classes=1000)

>>> # With customization
>>> model = (Template.resnet(layers=50)
...     .replace_activation('gelu')
...     .add_attention('se')
...     .build(num_classes=100))

class Template(config: Dict[str, Any])[source]

Bases: object

Base class for architecture templates with customization.

Templates allow you to start from a known architecture and customize it before building.

Example

>>> template = Template.resnet(layers=50)
>>> template.replace_activation('gelu')
>>> template.add_attention('se', after='conv')
>>> model = template.build(num_classes=100)

add_attention(attention_type: str = 'se', after: str = 'block') → Template[source]

Add attention mechanism.

Parameters:

attention_type – Type of attention (‘se’, ‘cbam’, ‘channel’, ‘spatial’)
after – Where to add (‘block’, ‘conv’, ‘stage’)

Returns:

Self for chaining

build(num_classes: int = 1000, **kwargs) → Module[source]

Build the model from template.

Parameters:

num_classes – Number of output classes
**kwargs – Additional build arguments

Returns:

Built nn.Module

classmethod densenet(layers: int = 121, **kwargs) → Template[source]: Create a DenseNet template.

describe() → str[source]: Get a human-readable description of the template.

classmethod from_name(name: str, **kwargs) → Template[source]

Create a template by name.

Parameters:

name – Template name (e.g., ‘resnet’, ‘vgg’)
**kwargs – Template-specific arguments

Returns:

Template instance

get_config() → Dict[str, Any][source]: Get the current configuration.

classmethod list_templates() → List[str][source]: List available template names.

classmethod mobilenet(version: int = 2, **kwargs) → Template[source]: Create a MobileNet template.

modify_stage(stage_idx: int, **kwargs) → Template[source]

Modify a specific stage.

Parameters:

stage_idx – Stage index (0-based)
**kwargs – Stage modifications (blocks, channels, etc.)

Returns:

Self for chaining

classmethod register(name: str)[source]: Decorator to register a template class.

replace_activation(activation: str) → Template[source]

Replace all activation functions.

Parameters:: activation – New activation function name
Returns:: Self for chaining

replace_norm(norm_type: str) → Template[source]

Replace all normalization layers.

Parameters:: norm_type – New normalization type (‘batch’, ‘layer’, ‘group’, ‘instance’)
Returns:: Self for chaining

classmethod resnet(layers: int = 50, **kwargs) → Template[source]: Create a ResNet template.

scale_channels(factor: float) → Template[source]

Scale all channel counts by a factor.

Parameters:: factor – Scaling factor (e.g., 0.5 for half, 2.0 for double)
Returns:: Self for chaining

set_head(**kwargs) → Template[source]

Configure the head.

Parameters:: **kwargs – Head configuration (hidden, dropout, etc.)
Returns:: Self for chaining

set_stem(**kwargs) → Template[source]

Configure the stem.

Parameters:: **kwargs – Stem configuration (channels, kernel, stride, etc.)
Returns:: Self for chaining

use_dropout(rate: float) → Template[source]

Set dropout rate.

Parameters:: rate – Dropout probability
Returns:: Self for chaining

classmethod vgg(layers: int = 16, **kwargs) → Template[source]: Create a VGG template.

resnet(layers: int = 50, num_classes: int = 1000, **kwargs) → Module[source]

Create a ResNet model.

Parametric factory function for ResNet architectures.

Parameters:

layers – Number of layers (18, 34, 50, 101, 152)
num_classes – Number of output classes
**kwargs – Additional configuration (activation, norm_type, use_se, etc.)

Returns:

ResNet model

Example

>>> model = resnet(layers=18, num_classes=10)
>>> model = resnet(layers=50, num_classes=1000, use_se=True)
>>> model = resnet(layers=101, activation='gelu')

vgg(layers: int = 16, num_classes: int = 1000, **kwargs) → Module[source]

Create a VGG model.

Parametric factory function for VGG architectures.

Parameters:

layers – Number of layers (11, 13, 16, 19)
num_classes – Number of output classes
**kwargs – Additional configuration (activation, use_batchnorm, dropout, etc.)

Returns:

VGG model

Example

>>> model = vgg(layers=16, num_classes=10)
>>> model = vgg(layers=19, num_classes=1000, use_batchnorm=True)

mobilenet(version: int = 2, num_classes: int = 1000, **kwargs) → Module[source]

Create a MobileNet model.

Parameters:

version – MobileNet version (1, 2, or 3)
num_classes – Number of output classes
**kwargs – Additional configuration (width_multiplier, mode, etc.)

Returns:

MobileNet model

Example

>>> model = mobilenet(version=1, num_classes=10)
>>> model = mobilenet(version=2, width_multiplier=0.5)
>>> model = mobilenet(version=3, mode='small')

densenet(layers: int = 121, num_classes: int = 1000, **kwargs) → Module[source]

Create a DenseNet model.

Parameters:

layers – Number of layers (121, 169, 201, 264)
num_classes – Number of output classes
**kwargs – Additional configuration (growth_rate, compression, etc.)

Returns:

DenseNet model

Example

>>> model = densenet(layers=121, num_classes=10)
>>> model = densenet(layers=201, growth_rate=48)

efficientnet(variant: str = 'b0', num_classes: int = 1000, **kwargs) → Module[source]

Create an EfficientNet model.

Parameters:

variant – Model variant (‘b0’ to ‘b7’)
num_classes – Number of output classes
**kwargs – Additional configuration

Returns:

EfficientNet model

Example

>>> model = efficientnet(variant='b0', num_classes=10)
>>> model = efficientnet(variant='b4', num_classes=1000)

class ResNetTemplate(layers: int = 50, **kwargs)[source]

Bases: Template

ResNet architecture template.

Example

>>> template = ResNetTemplate(layers=50)
>>> template.add_attention('se')
>>> model = template.build(num_classes=100)

add_attention(attention_type: str = 'se', after: str = 'block') → Template

Add attention mechanism.

Parameters:

attention_type – Type of attention (‘se’, ‘cbam’, ‘channel’, ‘spatial’)
after – Where to add (‘block’, ‘conv’, ‘stage’)

Returns:

Self for chaining

build(num_classes: int = 1000, **kwargs) → Module[source]: Build ResNet model from template configuration.

classmethod densenet(layers: int = 121, **kwargs) → Template: Create a DenseNet template.

describe() → str: Get a human-readable description of the template.

classmethod from_name(name: str, **kwargs) → Template

Create a template by name.

Parameters:

name – Template name (e.g., ‘resnet’, ‘vgg’)
**kwargs – Template-specific arguments

Returns:

Template instance

get_config() → Dict[str, Any]: Get the current configuration.

classmethod list_templates() → List[str]: List available template names.

classmethod mobilenet(version: int = 2, **kwargs) → Template: Create a MobileNet template.

modify_stage(stage_idx: int, **kwargs) → Template

Modify a specific stage.

Parameters:

stage_idx – Stage index (0-based)
**kwargs – Stage modifications (blocks, channels, etc.)

Returns:

Self for chaining

classmethod register(name: str): Decorator to register a template class.

replace_activation(activation: str) → Template

Replace all activation functions.

Parameters:: activation – New activation function name
Returns:: Self for chaining

replace_norm(norm_type: str) → Template

Replace all normalization layers.

Parameters:: norm_type – New normalization type (‘batch’, ‘layer’, ‘group’, ‘instance’)
Returns:: Self for chaining

classmethod resnet(layers: int = 50, **kwargs) → Template: Create a ResNet template.

scale_channels(factor: float) → Template

Scale all channel counts by a factor.

Parameters:: factor – Scaling factor (e.g., 0.5 for half, 2.0 for double)
Returns:: Self for chaining

set_head(**kwargs) → Template

Configure the head.

Parameters:: **kwargs – Head configuration (hidden, dropout, etc.)
Returns:: Self for chaining

set_stem(**kwargs) → Template

Configure the stem.

Parameters:: **kwargs – Stem configuration (channels, kernel, stride, etc.)
Returns:: Self for chaining

use_dropout(rate: float) → Template

Set dropout rate.

Parameters:: rate – Dropout probability
Returns:: Self for chaining

classmethod vgg(layers: int = 16, **kwargs) → Template: Create a VGG template.

class VGGTemplate(layers: int = 16, **kwargs)[source]

Bases: Template

VGG architecture template.

Example

>>> template = VGGTemplate(layers=16)
>>> template.use_batchnorm(True)
>>> model = template.build(num_classes=100)

add_attention(attention_type: str = 'se', after: str = 'block') → Template

Add attention mechanism.

Parameters:

attention_type – Type of attention (‘se’, ‘cbam’, ‘channel’, ‘spatial’)
after – Where to add (‘block’, ‘conv’, ‘stage’)

Returns:

Self for chaining

build(num_classes: int = 1000, **kwargs) → Module[source]: Build VGG model from template configuration.

classmethod densenet(layers: int = 121, **kwargs) → Template: Create a DenseNet template.

describe() → str: Get a human-readable description of the template.

classmethod from_name(name: str, **kwargs) → Template

Create a template by name.

Parameters:

name – Template name (e.g., ‘resnet’, ‘vgg’)
**kwargs – Template-specific arguments

Returns:

Template instance

get_config() → Dict[str, Any]: Get the current configuration.

classmethod list_templates() → List[str]: List available template names.

classmethod mobilenet(version: int = 2, **kwargs) → Template: Create a MobileNet template.

modify_stage(stage_idx: int, **kwargs) → Template

Modify a specific stage.

Parameters:

stage_idx – Stage index (0-based)
**kwargs – Stage modifications (blocks, channels, etc.)

Returns:

Self for chaining

classmethod register(name: str): Decorator to register a template class.

replace_activation(activation: str) → Template

Replace all activation functions.

Parameters:: activation – New activation function name
Returns:: Self for chaining

replace_norm(norm_type: str) → Template

Replace all normalization layers.

Parameters:: norm_type – New normalization type (‘batch’, ‘layer’, ‘group’, ‘instance’)
Returns:: Self for chaining

classmethod resnet(layers: int = 50, **kwargs) → Template: Create a ResNet template.

scale_channels(factor: float) → Template

Scale all channel counts by a factor.

Parameters:: factor – Scaling factor (e.g., 0.5 for half, 2.0 for double)
Returns:: Self for chaining

set_head(**kwargs) → Template

Configure the head.

Parameters:: **kwargs – Head configuration (hidden, dropout, etc.)
Returns:: Self for chaining

set_stem(**kwargs) → Template

Configure the stem.

Parameters:: **kwargs – Stem configuration (channels, kernel, stride, etc.)
Returns:: Self for chaining

use_batchnorm(enabled: bool = True) → VGGTemplate[source]: Enable or disable batch normalization.

use_dropout(rate: float) → Template

Set dropout rate.

Parameters:: rate – Dropout probability
Returns:: Self for chaining

classmethod vgg(layers: int = 16, **kwargs) → Template: Create a VGG template.

class MobileNetTemplate(version: int = 2, **kwargs)[source]

Bases: Template

MobileNet architecture template.

add_attention(attention_type: str = 'se', after: str = 'block') → Template

Add attention mechanism.

Parameters:

attention_type – Type of attention (‘se’, ‘cbam’, ‘channel’, ‘spatial’)
after – Where to add (‘block’, ‘conv’, ‘stage’)

Returns:

Self for chaining

build(num_classes: int = 1000, **kwargs) → Module[source]

Build the model from template.

Parameters:

num_classes – Number of output classes
**kwargs – Additional build arguments

Returns:

Built nn.Module

classmethod densenet(layers: int = 121, **kwargs) → Template: Create a DenseNet template.

describe() → str: Get a human-readable description of the template.

classmethod from_name(name: str, **kwargs) → Template

Create a template by name.

Parameters:

name – Template name (e.g., ‘resnet’, ‘vgg’)
**kwargs – Template-specific arguments

Returns:

Template instance

get_config() → Dict[str, Any]: Get the current configuration.

classmethod list_templates() → List[str]: List available template names.

classmethod mobilenet(version: int = 2, **kwargs) → Template: Create a MobileNet template.

modify_stage(stage_idx: int, **kwargs) → Template

Modify a specific stage.

Parameters:

stage_idx – Stage index (0-based)
**kwargs – Stage modifications (blocks, channels, etc.)

Returns:

Self for chaining

classmethod register(name: str): Decorator to register a template class.

replace_activation(activation: str) → Template

Replace all activation functions.

Parameters:: activation – New activation function name
Returns:: Self for chaining

replace_norm(norm_type: str) → Template

Replace all normalization layers.

Parameters:: norm_type – New normalization type (‘batch’, ‘layer’, ‘group’, ‘instance’)
Returns:: Self for chaining

classmethod resnet(layers: int = 50, **kwargs) → Template: Create a ResNet template.

scale_channels(factor: float) → Template

Scale all channel counts by a factor.

Parameters:: factor – Scaling factor (e.g., 0.5 for half, 2.0 for double)
Returns:: Self for chaining

set_head(**kwargs) → Template

Configure the head.

Parameters:: **kwargs – Head configuration (hidden, dropout, etc.)
Returns:: Self for chaining

set_stem(**kwargs) → Template

Configure the stem.

Parameters:: **kwargs – Stem configuration (channels, kernel, stride, etc.)
Returns:: Self for chaining

use_dropout(rate: float) → Template

Set dropout rate.

Parameters:: rate – Dropout probability
Returns:: Self for chaining

classmethod vgg(layers: int = 16, **kwargs) → Template: Create a VGG template.

class DenseNetTemplate(layers: int = 121, **kwargs)[source]

Bases: Template

DenseNet architecture template.

add_attention(attention_type: str = 'se', after: str = 'block') → Template

Add attention mechanism.

Parameters:

attention_type – Type of attention (‘se’, ‘cbam’, ‘channel’, ‘spatial’)
after – Where to add (‘block’, ‘conv’, ‘stage’)

Returns:

Self for chaining

build(num_classes: int = 1000, **kwargs) → Module[source]

Build the model from template.

Parameters:

num_classes – Number of output classes
**kwargs – Additional build arguments

Returns:

Built nn.Module

classmethod densenet(layers: int = 121, **kwargs) → Template: Create a DenseNet template.

describe() → str: Get a human-readable description of the template.

classmethod from_name(name: str, **kwargs) → Template

Create a template by name.

Parameters:

name – Template name (e.g., ‘resnet’, ‘vgg’)
**kwargs – Template-specific arguments

Returns:

Template instance

get_config() → Dict[str, Any]: Get the current configuration.

classmethod list_templates() → List[str]: List available template names.

classmethod mobilenet(version: int = 2, **kwargs) → Template: Create a MobileNet template.

modify_stage(stage_idx: int, **kwargs) → Template

Modify a specific stage.

Parameters:

stage_idx – Stage index (0-based)
**kwargs – Stage modifications (blocks, channels, etc.)

Returns:

Self for chaining

classmethod register(name: str): Decorator to register a template class.

replace_activation(activation: str) → Template

Replace all activation functions.

Parameters:: activation – New activation function name
Returns:: Self for chaining

replace_norm(norm_type: str) → Template

Replace all normalization layers.

Parameters:: norm_type – New normalization type (‘batch’, ‘layer’, ‘group’, ‘instance’)
Returns:: Self for chaining

classmethod resnet(layers: int = 50, **kwargs) → Template: Create a ResNet template.

scale_channels(factor: float) → Template

Scale all channel counts by a factor.

Parameters:: factor – Scaling factor (e.g., 0.5 for half, 2.0 for double)
Returns:: Self for chaining

set_head(**kwargs) → Template

Configure the head.

Parameters:: **kwargs – Head configuration (hidden, dropout, etc.)
Returns:: Self for chaining

set_stem(**kwargs) → Template

Configure the stem.

Parameters:: **kwargs – Stem configuration (channels, kernel, stride, etc.)
Returns:: Self for chaining

use_dropout(rate: float) → Template

Set dropout rate.

Parameters:: rate – Dropout probability
Returns:: Self for chaining

classmethod vgg(layers: int = 16, **kwargs) → Template: Create a VGG template.

class EfficientNetTemplate(variant: str = 'b0', **kwargs)[source]

Bases: Template

EfficientNet architecture template.

add_attention(attention_type: str = 'se', after: str = 'block') → Template

Add attention mechanism.

Parameters:

attention_type – Type of attention (‘se’, ‘cbam’, ‘channel’, ‘spatial’)
after – Where to add (‘block’, ‘conv’, ‘stage’)

Returns:

Self for chaining

build(num_classes: int = 1000, **kwargs) → Module[source]

Build the model from template.

Parameters:

num_classes – Number of output classes
**kwargs – Additional build arguments

Returns:

Built nn.Module

classmethod densenet(layers: int = 121, **kwargs) → Template: Create a DenseNet template.

describe() → str: Get a human-readable description of the template.

classmethod from_name(name: str, **kwargs) → Template

Create a template by name.

Parameters:

name – Template name (e.g., ‘resnet’, ‘vgg’)
**kwargs – Template-specific arguments

Returns:

Template instance

get_config() → Dict[str, Any]: Get the current configuration.

classmethod list_templates() → List[str]: List available template names.

classmethod mobilenet(version: int = 2, **kwargs) → Template: Create a MobileNet template.

modify_stage(stage_idx: int, **kwargs) → Template

Modify a specific stage.

Parameters:

stage_idx – Stage index (0-based)
**kwargs – Stage modifications (blocks, channels, etc.)

Returns:

Self for chaining

classmethod register(name: str): Decorator to register a template class.

replace_activation(activation: str) → Template

Replace all activation functions.

Parameters:: activation – New activation function name
Returns:: Self for chaining

replace_norm(norm_type: str) → Template

Replace all normalization layers.

Parameters:: norm_type – New normalization type (‘batch’, ‘layer’, ‘group’, ‘instance’)
Returns:: Self for chaining

classmethod resnet(layers: int = 50, **kwargs) → Template: Create a ResNet template.

scale_channels(factor: float) → Template

Scale all channel counts by a factor.

Parameters:: factor – Scaling factor (e.g., 0.5 for half, 2.0 for double)
Returns:: Self for chaining

set_head(**kwargs) → Template

Configure the head.

Parameters:: **kwargs – Head configuration (hidden, dropout, etc.)
Returns:: Self for chaining

set_stem(**kwargs) → Template

Configure the stem.

Parameters:: **kwargs – Stem configuration (channels, kernel, stride, etc.)
Returns:: Self for chaining

use_dropout(rate: float) → Template

Set dropout rate.

Parameters:: rate – Dropout probability
Returns:: Self for chaining

classmethod vgg(layers: int = 16, **kwargs) → Template: Create a VGG template.

Template Builders

resnet(layers: int = 50, num_classes: int = 1000, **kwargs) → Module[source]

Create a ResNet model.

Parametric factory function for ResNet architectures.

Parameters:

layers – Number of layers (18, 34, 50, 101, 152)
num_classes – Number of output classes
**kwargs – Additional configuration (activation, norm_type, use_se, etc.)

Returns:

ResNet model

Example

>>> model = resnet(layers=18, num_classes=10)
>>> model = resnet(layers=50, num_classes=1000, use_se=True)
>>> model = resnet(layers=101, activation='gelu')

vgg(layers: int = 16, num_classes: int = 1000, **kwargs) → Module[source]

Create a VGG model.

Parametric factory function for VGG architectures.

Parameters:

layers – Number of layers (11, 13, 16, 19)
num_classes – Number of output classes
**kwargs – Additional configuration (activation, use_batchnorm, dropout, etc.)

Returns:

VGG model

Example

>>> model = vgg(layers=16, num_classes=10)
>>> model = vgg(layers=19, num_classes=1000, use_batchnorm=True)

mobilenet(version: int = 2, num_classes: int = 1000, **kwargs) → Module[source]

Create a MobileNet model.

Parameters:

version – MobileNet version (1, 2, or 3)
num_classes – Number of output classes
**kwargs – Additional configuration (width_multiplier, mode, etc.)

Returns:

MobileNet model

Example

>>> model = mobilenet(version=1, num_classes=10)
>>> model = mobilenet(version=2, width_multiplier=0.5)
>>> model = mobilenet(version=3, mode='small')

densenet(layers: int = 121, num_classes: int = 1000, **kwargs) → Module[source]

Create a DenseNet model.

Parameters:

layers – Number of layers (121, 169, 201, 264)
num_classes – Number of output classes
**kwargs – Additional configuration (growth_rate, compression, etc.)

Returns:

DenseNet model

Example

>>> model = densenet(layers=121, num_classes=10)
>>> model = densenet(layers=201, growth_rate=48)

efficientnet(variant: str = 'b0', num_classes: int = 1000, **kwargs) → Module[source]

Create an EfficientNet model.

Parameters:

variant – Model variant (‘b0’ to ‘b7’)
num_classes – Number of output classes
**kwargs – Additional configuration

Returns:

EfficientNet model

Example

>>> model = efficientnet(variant='b0', num_classes=10)
>>> model = efficientnet(variant='b4', num_classes=1000)

Base Classes

class Template(config: Dict[str, Any])[source]

Bases: object

Base class for architecture templates with customization.

Templates allow you to start from a known architecture and customize it before building.

Example

>>> template = Template.resnet(layers=50)
>>> template.replace_activation('gelu')
>>> template.add_attention('se', after='conv')
>>> model = template.build(num_classes=100)

classmethod register(name: str)[source]: Decorator to register a template class.

classmethod from_name(name: str, **kwargs) → Template[source]

Create a template by name.

Parameters:

name – Template name (e.g., ‘resnet’, ‘vgg’)
**kwargs – Template-specific arguments

Returns:

Template instance

classmethod resnet(layers: int = 50, **kwargs) → Template[source]: Create a ResNet template.

classmethod vgg(layers: int = 16, **kwargs) → Template[source]: Create a VGG template.

classmethod mobilenet(version: int = 2, **kwargs) → Template[source]: Create a MobileNet template.

classmethod densenet(layers: int = 121, **kwargs) → Template[source]: Create a DenseNet template.

classmethod list_templates() → List[str][source]: List available template names.

replace_activation(activation: str) → Template[source]

Replace all activation functions.

Parameters:: activation – New activation function name
Returns:: Self for chaining

replace_norm(norm_type: str) → Template[source]

Replace all normalization layers.

Parameters:: norm_type – New normalization type (‘batch’, ‘layer’, ‘group’, ‘instance’)
Returns:: Self for chaining

add_attention(attention_type: str = 'se', after: str = 'block') → Template[source]

Add attention mechanism.

Parameters:

attention_type – Type of attention (‘se’, ‘cbam’, ‘channel’, ‘spatial’)
after – Where to add (‘block’, ‘conv’, ‘stage’)

Returns:

Self for chaining

modify_stage(stage_idx: int, **kwargs) → Template[source]

Modify a specific stage.

Parameters:

stage_idx – Stage index (0-based)
**kwargs – Stage modifications (blocks, channels, etc.)

Returns:

Self for chaining

set_stem(**kwargs) → Template[source]

Configure the stem.

Parameters:: **kwargs – Stem configuration (channels, kernel, stride, etc.)
Returns:: Self for chaining

set_head(**kwargs) → Template[source]

Configure the head.

Parameters:: **kwargs – Head configuration (hidden, dropout, etc.)
Returns:: Self for chaining

use_dropout(rate: float) → Template[source]

Set dropout rate.

Parameters:: rate – Dropout probability
Returns:: Self for chaining

scale_channels(factor: float) → Template[source]

Scale all channel counts by a factor.

Parameters:: factor – Scaling factor (e.g., 0.5 for half, 2.0 for double)
Returns:: Self for chaining

get_config() → Dict[str, Any][source]: Get the current configuration.

describe() → str[source]: Get a human-readable description of the template.

build(num_classes: int = 1000, **kwargs) → Module[source]

Build the model from template.

Parameters:

num_classes – Number of output classes
**kwargs – Additional build arguments

Returns:

Built nn.Module