Blocks Module

The blocks module provides the fundamental neural network building blocks used in PyNAS architecture search. These modules include various convolution types, activation functions, pooling layers, and specialized heads for different tasks.

Overview

The blocks module is organized into several submodules:

• Convolutions: Various convolution blocks including classic, mobile, residual, and dense variants

• Activations: Activation functions like ReLU, GELU, Sigmoid, and Softmax

• Pooling: Max and average pooling operations

• Heads: Classification and segmentation heads for different tasks

• Residual: Residual connection utilities and implementations

Convolution Blocks

class pynas.blocks.convolutions.ConvAct(in_channels, out_channels, kernel_size=3, stride=1, padding=1, activation=ReLU)

Bases: Sequential

__init__(in_channels, out_channels, kernel_size=3, stride=1, padding=1, activation=ReLU)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

class pynas.blocks.convolutions.ConvBnAct(in_channels, out_channels, kernel_size=3, stride=1, padding=1, activation=ReLU)

Bases: Sequential

__init__(in_channels, out_channels, kernel_size=3, stride=1, padding=1, activation=ReLU)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

class pynas.blocks.convolutions.SEBlock(in_channels, reduction=16)

Bases: Module

__init__(in_channels, reduction=16)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

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.

class pynas.blocks.convolutions.ConvSE(in_channels, out_channels, kernel_size=3, stride=1, padding=1, activation=ReLU)

Bases: Sequential

__init__(in_channels, out_channels, kernel_size=3, stride=1, padding=1, activation=ReLU)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

class pynas.blocks.convolutions.MBConv(in_channels, out_channels, dw_kernel_size=3, expansion_factor=4, activation=ReLU)

Bases: Module

__init__(in_channels, out_channels, dw_kernel_size=3, expansion_factor=4, activation=ReLU)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

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.

class pynas.blocks.convolutions.MBConvNoRes(in_channels, out_channels, dw_kernel_size=3, expansion_factor=4, activation=ReLU)

Bases: Module

__init__(in_channels, out_channels, dw_kernel_size=3, expansion_factor=4, activation=ReLU)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

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.

class pynas.blocks.convolutions.CSPConvBlock(in_channels, num_blocks=1, activation=ReLU)

Bases: Module

__init__(in_channels, num_blocks=1, activation=ReLU)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

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.

class pynas.blocks.convolutions.CSPMBConvBlock(in_channels, num_blocks=1, dw_kernel_size=3, expansion_factor=4, activation=ReLU)

Bases: Module

__init__(in_channels, num_blocks=1, dw_kernel_size=3, expansion_factor=4, activation=ReLU)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

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.

class pynas.blocks.convolutions.DenseNetBlock(in_channels, out_channels, activation=ReLU)

Bases: Module

Basic DenseNet block composed by one 3x3 convs with residual connection. The residual connection is perfomed by concatenate the input and the output.

__init__(in_channels, out_channels, activation=ReLU)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

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.

class pynas.blocks.convolutions.ResNetBasicBlock(in_channels, out_channels, reduction_factor=4, activation=ReLU)

Bases: Module

__init__(in_channels, out_channels, reduction_factor=4, activation=ReLU)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

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.

class pynas.blocks.convolutions.ResNetBlock(in_channels, out_channels, reduction_factor=4, activation=ReLU)

Bases: Module

__init__(in_channels, out_channels, reduction_factor=4, activation=ReLU)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

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.

class pynas.blocks.convolutions.Upsample(scale_factor=2, mode='nearest')

Bases: Module

__init__(scale_factor=2, mode='nearest')

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

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.

Classic Convolution Blocks

class pynas.blocks.convolutions.ConvAct(in_channels, out_channels, kernel_size=3, stride=1, padding=1, activation=ReLU)

Bases: Sequential

Basic convolution followed by activation.

__init__(in_channels, out_channels, kernel_size=3, stride=1, padding=1, activation=ReLU)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

class pynas.blocks.convolutions.ConvBnAct(in_channels, out_channels, kernel_size=3, stride=1, padding=1, activation=ReLU)

Bases: Sequential

Convolution, batch normalization, and activation in sequence.

__init__(in_channels, out_channels, kernel_size=3, stride=1, padding=1, activation=ReLU)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

class pynas.blocks.convolutions.ConvSE(in_channels, out_channels, kernel_size=3, stride=1, padding=1, activation=ReLU)

Bases: Sequential

Convolution with Squeeze-and-Excitation attention mechanism.

__init__(in_channels, out_channels, kernel_size=3, stride=1, padding=1, activation=ReLU)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

Squeeze-and-Excitation Block

class pynas.blocks.convolutions.SEBlock(in_channels, reduction=16)

Bases: Module

Squeeze-and-Excitation block for channel attention.

__init__(in_channels, reduction=16)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

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.

Mobile Convolution Blocks

class pynas.blocks.convolutions.MBConv(in_channels, out_channels, dw_kernel_size=3, expansion_factor=4, activation=ReLU)

Bases: Module

Mobile Inverted Bottleneck Convolution with residual connection.

__init__(in_channels, out_channels, dw_kernel_size=3, expansion_factor=4, activation=ReLU)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

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.

class pynas.blocks.convolutions.MBConvNoRes(in_channels, out_channels, dw_kernel_size=3, expansion_factor=4, activation=ReLU)

Bases: Module

Mobile Inverted Bottleneck Convolution without residual connection.

__init__(in_channels, out_channels, dw_kernel_size=3, expansion_factor=4, activation=ReLU)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

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.

CSP Blocks

class pynas.blocks.convolutions.CSPConvBlock(in_channels, num_blocks=1, activation=ReLU)

Bases: Module

Cross Stage Partial convolution block for efficient feature learning.

__init__(in_channels, num_blocks=1, activation=ReLU)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

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.

class pynas.blocks.convolutions.CSPMBConvBlock(in_channels, num_blocks=1, dw_kernel_size=3, expansion_factor=4, activation=ReLU)

Bases: Module

CSP block using Mobile Inverted Bottleneck convolutions.

__init__(in_channels, num_blocks=1, dw_kernel_size=3, expansion_factor=4, activation=ReLU)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

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.

Dense and Residual Blocks

class pynas.blocks.convolutions.DenseNetBlock(in_channels, out_channels, activation=ReLU)

Bases: Module

Basic DenseNet block composed by one 3x3 convs with residual connection. The residual connection is perfomed by concatenate the input and the output.

DenseNet block with concatenation-based feature reuse.

__init__(in_channels, out_channels, activation=ReLU)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

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.

class pynas.blocks.convolutions.ResNetBasicBlock(in_channels, out_channels, reduction_factor=4, activation=ReLU)

Bases: Module

Basic ResNet block with bottleneck structure.

__init__(in_channels, out_channels, reduction_factor=4, activation=ReLU)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

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.

class pynas.blocks.convolutions.ResNetBlock(in_channels, out_channels, reduction_factor=4, activation=ReLU)

Bases: Module

ResNet block with residual connection.

__init__(in_channels, out_channels, reduction_factor=4, activation=ReLU)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

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.

Utility Blocks

class pynas.blocks.convolutions.Upsample(scale_factor=2, mode='nearest')

Bases: Module

Upsampling layer for increasing spatial resolution.

__init__(scale_factor=2, mode='nearest')

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

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.

Activation Functions

class pynas.blocks.activations.GELU(*args, **kwargs)

Bases: Module

Gaussian Error Linear Unit (GELU) activation function. This module implements the GELU activation function, which is used to introduce non-linearity in the network. GELU is a smooth, non-monotonic function that models the Gaussian cumulative distribution function. It is commonly used in transformer architectures and other advanced models.

Parameters:

x (Tensor) – Input tensor to which the GELU activation function is applied.

Returns:

Output tensor after applying the GELU activation function.

Return type:

Tensor

forward(x)

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.

class pynas.blocks.activations.ReLU(*args, **kwargs)

Bases: Module

Rectified Linear Unit (ReLU) activation function. This module implements the ReLU activation function, which is widely used in neural networks for introducing non-linearity. ReLU is defined as the positive part of its argument, where each element of the input tensor x that is less than zero is replaced with zero. This function increases the non-linear properties of the decision function and the overall network without affecting the receptive fields of the convolution layer.

Parameters:

x (Tensor) – Input tensor to which the ReLU activation function is applied.

Returns:

Output tensor after applying the ReLU activation function.

Return type:

Tensor

forward(x)

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.

class pynas.blocks.activations.Sigmoid(*args, **kwargs)

Bases: Module

Sigmoid activation function.

forward(x)

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.

class pynas.blocks.activations.Softmax(*args, **kwargs)

Bases: Module

Softmax activation function.

forward(x)

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.

class pynas.blocks.activations.GELU(*args, **kwargs)

Bases: Module

Gaussian Error Linear Unit (GELU) activation function. This module implements the GELU activation function, which is used to introduce non-linearity in the network. GELU is a smooth, non-monotonic function that models the Gaussian cumulative distribution function. It is commonly used in transformer architectures and other advanced models.

Parameters:

x (Tensor) – Input tensor to which the GELU activation function is applied.

Returns:

Output tensor after applying the GELU activation function.

Return type:

Tensor

Gaussian Error Linear Unit activation function.

forward(x)

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.

class pynas.blocks.activations.ReLU(*args, **kwargs)

Bases: Module

Rectified Linear Unit (ReLU) activation function. This module implements the ReLU activation function, which is widely used in neural networks for introducing non-linearity. ReLU is defined as the positive part of its argument, where each element of the input tensor x that is less than zero is replaced with zero. This function increases the non-linear properties of the decision function and the overall network without affecting the receptive fields of the convolution layer.

Parameters:

x (Tensor) – Input tensor to which the ReLU activation function is applied.

Returns:

Output tensor after applying the ReLU activation function.

Return type:

Tensor

Rectified Linear Unit activation function.

forward(x)

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.

class pynas.blocks.activations.Sigmoid(*args, **kwargs)

Bases: Module

Sigmoid activation function.

Sigmoid activation function.

forward(x)

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.

class pynas.blocks.activations.Softmax(*args, **kwargs)

Bases: Module

Softmax activation function.

Softmax activation function for multi-class classification.

forward(x)

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.

Pooling Operations

class pynas.blocks.pooling.AvgPool(kernel_size=2, stride=2)

Bases: Sequential

Implements an average pooling layer. This layer applies a 2D average pooling over an input signal (usually an image) represented as a batch of multichannel data. It reduces the spatial dimensions (width and height) of the input by taking the average of elements in a kernel-sized window, which slides over the input data with a specified stride.

Parameters: - kernel_size (int, optional): The size of the window for each dimension of the input tensor. Default is 2. - stride (int, optional): The stride of the window. Default is 2.

Example Usage:

avg_pool_layer = AvgPool(kernel_size=2, stride=2)

__init__(kernel_size=2, stride=2)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

class pynas.blocks.pooling.MaxPool(kernel_size=2, stride=None, padding=0)

Bases: Sequential

Implements a max pooling layer. This layer applies a 2D max pooling over an input signal (usually an image) represented as a batch of multichannel data. It reduces the spatial dimensions (width and height) of the input by taking the maximum value of elements in a kernel-sized window, which slides over the input data with a specified stride and padding.

Parameters: - kernel_size (int): The size of the window for each dimension of the input tensor. - stride (int, optional): The stride of the window. Defaults to kernel_size if not specified. - padding (int, optional): The amount of padding added to all sides of the input. Default is 0.

Example Usage:

max_pool_layer = MaxPool(kernel_size=2, stride=2, padding=0)

__init__(kernel_size=2, stride=None, padding=0)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

class pynas.blocks.pooling.AvgPool(kernel_size=2, stride=2)

Bases: Sequential

Implements an average pooling layer. This layer applies a 2D average pooling over an input signal (usually an image) represented as a batch of multichannel data. It reduces the spatial dimensions (width and height) of the input by taking the average of elements in a kernel-sized window, which slides over the input data with a specified stride.

Parameters: - kernel_size (int, optional): The size of the window for each dimension of the input tensor. Default is 2. - stride (int, optional): The stride of the window. Default is 2.

Example Usage:

avg_pool_layer = AvgPool(kernel_size=2, stride=2)

Average pooling operation for spatial downsampling.

__init__(kernel_size=2, stride=2)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

class pynas.blocks.pooling.MaxPool(kernel_size=2, stride=None, padding=0)

Bases: Sequential

Implements a max pooling layer. This layer applies a 2D max pooling over an input signal (usually an image) represented as a batch of multichannel data. It reduces the spatial dimensions (width and height) of the input by taking the maximum value of elements in a kernel-sized window, which slides over the input data with a specified stride and padding.

Parameters: - kernel_size (int): The size of the window for each dimension of the input tensor. - stride (int, optional): The stride of the window. Defaults to kernel_size if not specified. - padding (int, optional): The amount of padding added to all sides of the input. Default is 0.

Example Usage:

max_pool_layer = MaxPool(kernel_size=2, stride=2, padding=0)

Max pooling operation for spatial downsampling.

__init__(kernel_size=2, stride=None, padding=0)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

Head Modules

class pynas.blocks.heads.Dropout(p=0.5, inplace=False)

Bases: Module

Dropout layer for regularization in neural networks.

Parameters:

• p (float, optional) – Probability of an element to be zeroed. Default: 0.5

• inplace (bool, optional) – If set to True, will do this operation in-place. Default: False

__init__(p=0.5, inplace=False)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

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.

class pynas.blocks.heads.MultiInputClassifier(input_shapes, common_dim=256, mlp_depth=2, mlp_hidden_dim=512, num_classes=10, use_adaptive_pool=True, pool_size=(4, 4))

Bases: Module

A PyTorch module for a multi-input classifier that processes multiple input tensors with different shapes and combines their features for classification. :param input_shapes: A list of shapes for each input tensor,

excluding the batch dimension. Each shape can be either (C, H, W) for spatial inputs or (D,) for flat vector inputs.

Parameters:

• common_dim (int, optional) – The dimension to which all inputs are projected. Defaults to 256.

• mlp_depth (int, optional) – The number of layers in the final MLP classifier. Defaults to 2.

• mlp_hidden_dim (int, optional) – The number of hidden units in each MLP layer. Defaults to 512.

• num_classes (int, optional) – The number of output classes for classification. Defaults to 10.

• use_adaptive_pool (bool, optional) – Whether to apply adaptive average pooling for spatial inputs. Defaults to True.

• pool_size (Tuple[int, int], optional) – The target size for adaptive pooling if it is used. Defaults to (4, 4).

• input_shapes (List[Tuple[int, ...]])

projections

A list of projection modules for each input tensor. These modules transform the inputs to the common dimension.

Type:

nn.ModuleList

flatten

A module to flatten the projected tensors.

Type:

nn.Flatten

total_input_dim

The total input dimension after concatenating all projected tensors.

Type:

int

classifier

The MLP classifier that processes the concatenated features and outputs class probabilities.

Type:

nn.Sequential

forward(inputs

List[torch.Tensor]) -> torch.Tensor: Processes the input tensors, projects them to a common dimension, concatenates their features, and passes them through the MLP classifier to produce the output logits.

Raises:

ValueError – If an input shape is not supported (e.g., not (C, H, W) or (D,)).

Parameters:

• input_shapes (List[Tuple[int, ...]])

• common_dim (int)

• mlp_depth (int)

• mlp_hidden_dim (int)

• num_classes (int)

• use_adaptive_pool (bool)

• pool_size (Tuple[int, int])

__init__(input_shapes, common_dim=256, mlp_depth=2, mlp_hidden_dim=512, num_classes=10, use_adaptive_pool=True, pool_size=(4, 4))

Initializes internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

• input_shapes (List[Tuple[int, ...]])

• common_dim (int)

• mlp_depth (int)

• mlp_hidden_dim (int)

• num_classes (int)

• use_adaptive_pool (bool)

• pool_size (Tuple[int, int])

forward(inputs)

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.

Parameters:

inputs (List[Tensor])

Return type:

Tensor

class pynas.blocks.heads.Classifier(encoder, dm, verbose=False)

Bases: object

__init__(encoder, dm, verbose=False)

build_head(input_shape=(1, 2, 256, 256))

dummy_test()

forward(x)

class pynas.blocks.heads.Dropout(p=0.5, inplace=False)

Bases: Module

Dropout layer for regularization in neural networks.

Parameters:

• p (float, optional) – Probability of an element to be zeroed. Default: 0.5

• inplace (bool, optional) – If set to True, will do this operation in-place. Default: False

Dropout layer for regularization.

__init__(p=0.5, inplace=False)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

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.

class pynas.blocks.heads.Classifier(encoder, dm, verbose=False)

Bases: object

Classification head that combines encoder with adaptive pooling and linear layers.

__init__(encoder, dm, verbose=False)

build_head(input_shape=(1, 2, 256, 256))

dummy_test()

forward(x)

class pynas.blocks.heads.MultiInputClassifier(input_shapes, common_dim=256, mlp_depth=2, mlp_hidden_dim=512, num_classes=10, use_adaptive_pool=True, pool_size=(4, 4))

Bases: Module

A PyTorch module for a multi-input classifier that processes multiple input tensors with different shapes and combines their features for classification. :param input_shapes: A list of shapes for each input tensor,

excluding the batch dimension. Each shape can be either (C, H, W) for spatial inputs or (D,) for flat vector inputs.

Parameters:

• common_dim (int, optional) – The dimension to which all inputs are projected. Defaults to 256.

• mlp_depth (int, optional) – The number of layers in the final MLP classifier. Defaults to 2.

• mlp_hidden_dim (int, optional) – The number of hidden units in each MLP layer. Defaults to 512.

• num_classes (int, optional) – The number of output classes for classification. Defaults to 10.

• use_adaptive_pool (bool, optional) – Whether to apply adaptive average pooling for spatial inputs. Defaults to True.

• pool_size (Tuple[int, int], optional) – The target size for adaptive pooling if it is used. Defaults to (4, 4).

• input_shapes (List[Tuple[int, ...]])

projections

A list of projection modules for each input tensor. These modules transform the inputs to the common dimension.

Type:

nn.ModuleList

flatten

A module to flatten the projected tensors.

Type:

nn.Flatten

total_input_dim

The total input dimension after concatenating all projected tensors.

Type:

int

classifier

The MLP classifier that processes the concatenated features and outputs class probabilities.

Type:

nn.Sequential

forward(inputs

List[torch.Tensor]) -> torch.Tensor: Processes the input tensors, projects them to a common dimension, concatenates their features, and passes them through the MLP classifier to produce the output logits.

Raises:

ValueError – If an input shape is not supported (e.g., not (C, H, W) or (D,)).

Parameters:

• input_shapes (List[Tuple[int, ...]])

• common_dim (int)

• mlp_depth (int)

• mlp_hidden_dim (int)

• num_classes (int)

• use_adaptive_pool (bool)

• pool_size (Tuple[int, int])

Multi-input classifier for handling multiple input tensors with different shapes.

__init__(input_shapes, common_dim=256, mlp_depth=2, mlp_hidden_dim=512, num_classes=10, use_adaptive_pool=True, pool_size=(4, 4))

Initializes internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

• input_shapes (List[Tuple[int, ...]])

• common_dim (int)

• mlp_depth (int)

• mlp_hidden_dim (int)

• num_classes (int)

• use_adaptive_pool (bool)

• pool_size (Tuple[int, int])

forward(inputs)

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.

Parameters:

inputs (List[Tensor])

Return type:

Tensor

Residual Utilities

class pynas.blocks.residual.Residual(block, res_func=None, shortcut=None, *args, **kwargs)

Bases: Module

It applies residual connection to a nn.Module where the output becomes

\(y = F(x) + x\)

Examples

>>> block = nn.Identity() // does nothing
>>> res = Residual(block, res_func=lambda x, res: x + res)
>>> res(x) // tensor([2])

You can also pass a shortcut function

>>> res = Residual(block, res_func=lambda x, res: x + res, shortcut=lambda x: x * 2)
>>> res(x) // tensor([3])

Parameters:

• block (Module)

• res_func (Callable[[Tensor, Tensor], Tensor])

• shortcut (Module)

__init__(block, res_func=None, shortcut=None, *args, **kwargs)

Parameters:

• block (nn.Module) – A Pytorch module

• res_func (Callable[[Tensor], Tensor], optional) – The residual function. Defaults to None.

• shortcut (nn.Module, optional) – A function applied before the input is passed to block. Defaults to None.

forward(x)

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.

Parameters:

x (Tensor)

Return type:

Tensor

pynas.blocks.residual.add(x, res)

Parameters:

• x (Tensor)

• res (Tensor)

Return type:

Tensor

class pynas.blocks.residual.ResidualAdd(*args, **kwags)

Bases: Residual

__init__(*args, **kwags)

Parameters:

• block (nn.Module) – A Pytorch module

• res_func (Callable[[Tensor], Tensor], optional) – The residual function. Defaults to None.

• shortcut (nn.Module, optional) – A function applied before the input is passed to block. Defaults to None.

class pynas.blocks.residual.InputForward(blocks, aggr_func)

Bases: Module

This module passes the input to multiple modules and applies a aggregation function on the result.

Parameters:

• blocks (Module)

• aggr_func (Callable[[Tensor], Tensor])

__init__(blocks, aggr_func)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

• blocks (Module)

• aggr_func (Callable[[Tensor], Tensor])

forward(x)

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.

Parameters:

x (Tensor)

Return type:

Tensor

pynas.blocks.residual.Cat2d = functools.partial(<class 'pynas.blocks.residual.InputForward'>, aggr_func=<function <lambda>>)

Pass the input to multiple modules and concatenates the output, for 1D input you can use Cat, while for 2D inputs, such as images, you can use Cat2d.

Examples

>>> blocks = nn.ModuleList([nn.Conv2d(32, 64, kernel_size=3), nn.Conv2d(32, 64, kernel_size=3)])
>>> x = torch.rand(1, 32, 48, 48)
>>> Cat2d(blocks)(x).shape
# torch.Size([1, 128, 46, 46])

Parameters:

• blocks (Module)

• aggr_func (Callable[[Tensor], Tensor])

class pynas.blocks.residual.Residual(block, res_func=None, shortcut=None, *args, **kwargs)

Bases: Module

It applies residual connection to a nn.Module where the output becomes

\(y = F(x) + x\)

Examples

>>> block = nn.Identity() // does nothing
>>> res = Residual(block, res_func=lambda x, res: x + res)
>>> res(x) // tensor([2])

You can also pass a shortcut function

>>> res = Residual(block, res_func=lambda x, res: x + res, shortcut=lambda x: x * 2)
>>> res(x) // tensor([3])

Generic residual connection wrapper for any PyTorch module.

Parameters:

• block (Module)

• res_func (Callable[[Tensor, Tensor], Tensor])

• shortcut (Module)

__init__(block, res_func=None, shortcut=None, *args, **kwargs)

Parameters:

• block (nn.Module) – A Pytorch module

• res_func (Callable[[Tensor], Tensor], optional) – The residual function. Defaults to None.

• shortcut (nn.Module, optional) – A function applied before the input is passed to block. Defaults to None.

forward(x)

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.

Parameters:

x (Tensor)

Return type:

Tensor

class pynas.blocks.residual.ResidualAdd(*args, **kwags)

Bases: Residual

Residual connection with element-wise addition.

__init__(*args, **kwags)

Parameters:

• block (nn.Module) – A Pytorch module

• res_func (Callable[[Tensor], Tensor], optional) – The residual function. Defaults to None.

• shortcut (nn.Module, optional) – A function applied before the input is passed to block. Defaults to None.

class pynas.blocks.residual.InputForward(blocks, aggr_func)

Bases: Module

This module passes the input to multiple modules and applies a aggregation function on the result.

Module that passes input through multiple modules and aggregates results.

Parameters:

• blocks (Module)

• aggr_func (Callable[[Tensor], Tensor])

__init__(blocks, aggr_func)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

• blocks (Module)

• aggr_func (Callable[[Tensor], Tensor])

forward(x)

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.

Parameters:

x (Tensor)

Return type:

Tensor

Utility Functions

pynas.blocks.residual.add(x, res)

Element-wise addition function for residual connections.

Parameters:

• x (Tensor)

• res (Tensor)

Return type:

Tensor

Lambda Module

class pynas.blocks.Lambda(lambd)

Bases: Module

An utility Module, it allows custom function to be passed

Parameters:

lambd (Callable[Tensor]) – A function that does something on a tensor

Examples

>>> add_two = Lambda(lambd x: x + 2)
>>> add_two(Tensor([0])) // 2

Utility module for wrapping custom functions as PyTorch modules.

__init__(lambd)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

lambd (Callable[[Tensor], Tensor])

forward(x)

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.

Parameters:

x (Tensor)

Return type:

Tensor

Block Configuration

Block Types and Parameters

The following table shows the available block types and their configurable parameters:

Block Types and Parameters

Block Type	Parameters
ConvAct	out_channels_coefficient, kernel_size, stride, padding, activation
ConvBnAct	out_channels_coefficient, kernel_size, stride, padding, activation
ConvSE	out_channels_coefficient, kernel_size, stride, padding, activation
MBConv	expansion_factor, dw_kernel_size, activation
MBConvNoRes	expansion_factor, dw_kernel_size, activation
CSPConvBlock	num_blocks, activation
CSPMBConvBlock	expansion_factor, dw_kernel_size, num_blocks, activation
DenseNetBlock	out_channels_coefficient, activation
ResNetBlock	reduction_factor, activation
AvgPool	kernel_size, stride
MaxPool	kernel_size, stride
Dropout	dropout_rate

Block Vocabulary

PyNAS uses a vocabulary system for encoding architectures:

Convolution Blocks: - b: ConvAct - c: ConvBnAct - e: ConvSE (with Squeeze-and-Excitation) - d: DenseNetBlock - m: MBConv - n: MBConvNoRes - R: ResNetBlock - D: Dropout

Activation Functions: - r: ReLU - g: GELU - sg: Sigmoid - s: Softmax

Pooling Operations: - a: AvgPool - M: MaxPool

Usage Examples

Basic Convolution Block

from pynas.blocks.convolutions import ConvBnAct
from pynas.blocks.activations import ReLU

# Create a convolution block
conv_block = ConvBnAct(
    in_channels=32,
    out_channels=64,
    kernel_size=3,
    stride=1,
    padding=1,
    activation=ReLU
)

# Use in forward pass
output = conv_block(input_tensor)

Mobile Inverted Bottleneck

from pynas.blocks.convolutions import MBConv
from pynas.blocks.activations import ReLU

# Create MBConv block
mb_block = MBConv(
    in_channels=64,
    out_channels=64,
    expansion_factor=4,
    dw_kernel_size=3,
    activation=ReLU
)

output = mb_block(input_tensor)

Classification Head

from pynas.blocks.heads import Classifier
import torch.nn as nn

# Assume we have a data module with num_classes and input_shape
class SimpleDataModule:
    def __init__(self):
        self.num_classes = 10
        self.input_shape = (3, 224, 224)

# Create encoder (any PyTorch model)
encoder = nn.Sequential(
    nn.Conv2d(3, 64, 3, padding=1),
    nn.ReLU(),
    nn.AdaptiveAvgPool2d((7, 7))
)

# Create classifier
dm = SimpleDataModule()
classifier = Classifier(encoder=encoder, dm=dm, verbose=True)

# The classifier combines encoder + classification head
output = classifier.model(input_tensor)

Residual Connections

from pynas.blocks.residual import Residual, ResidualAdd
from pynas.blocks.convolutions import ConvBnAct
from pynas.blocks.activations import ReLU

# Create a block with residual connection
conv_block = ConvBnAct(64, 64, activation=ReLU)
residual_block = ResidualAdd(conv_block)

# Or use custom residual function
custom_residual = Residual(
    block=conv_block,
    res_func=lambda x, res: x + res * 0.5  # Custom weighting
)

output = residual_block(input_tensor)

Architecture Encoding

# Example architecture string: "c3r" means 3 ConvBnAct layers with ReLU
# This gets parsed by the architecture builder into:

architecture = [
    {'layer_type': 'ConvBnAct', 'activation': 'ReLU'},
    {'layer_type': 'ConvBnAct', 'activation': 'ReLU'},
    {'layer_type': 'ConvBnAct', 'activation': 'ReLU'}
]

Best Practices

1. Channel Management: Use appropriate out_channels_coefficient values to control model complexity

2. Activation Choice: ReLU for general use, GELU for transformer-like architectures

3. Pooling Strategy: Use MaxPool for feature extraction, AvgPool before classification

4. Residual Connections: Essential for deep networks to avoid gradient vanishing

5. Mobile Blocks: Use MBConv for efficient mobile/edge deployments

6. CSP Blocks: Good balance between efficiency and representational power

Performance Considerations

• MBConv: Most parameter-efficient for mobile deployment

• ConvBnAct: Standard choice for balanced performance

• CSP Blocks: Good for avoiding gradient bottlenecks

• DenseNet Blocks: High memory usage but strong feature reuse

• ResNet Blocks: Stable training for very deep networks

The blocks in this module are designed to be modular and composable, allowing the genetic algorithm to efficiently explore different architectural combinations while maintaining compatibility across different block types.