zae_engine.models.converter package

Submodules

zae_engine.models.converter.dim_converter module

class zae_engine.models.converter.dim_converter.DimConverter(model: Module)

Bases: object

A class to convert the dimensionality of layers in a given PyTorch model from 1D to 2D, 2D to 3D, or vice versa.

This class identifies layers in a model that can be converted to different dimensions, such as Conv1d to Conv2d, and adapts their weights and parameters accordingly.

model

The original PyTorch model to be converted.

Type:

nn.Module

module_dict

A dictionary storing the original layers of the model.

Type:

dict

new_module_dict

A dictionary storing the converted layers of the model.

Type:

dict

layer_dict

A dictionary mapping layer names to their respective modules.

Type:

dict

param_dict

A dictionary mapping parameter names to their respective tensors.

Type:

dict

find_convertable(model: nn.Module) → tuple[dict, dict]:

Finds and returns layers and parameters in the model that can be converted.

dim_correction(reduce: bool):

Corrects the dimensionality of the layers identified for conversion.

const_getter(conv_module: nn.Module, reduce: bool) → dict:

Retrieves and adjusts the constants of a given convolutional module.

apply_new_dict(base_module: nn.Module, name: str, module: nn.Module):

Applies the converted layers to a new model structure.

convert(pattern: str, \*args, \*\*kwargs) → nn.Module:

Converts the dimensionality of the model based on the specified pattern.

apply_new_dict(base_module: Module, name: str, module: Module)

Apply the converted layers to a new model structure.

Parameters:

• base_module (nn.Module) – The base module to which the converted layers will be applied.

• name (str) – The name of the layer to be applied.

• module (nn.Module) – The converted module to be applied.

static const_getter(conv_module: Module, reduce: bool = None) → dict

Retrieve and adjust the constants of a given convolutional module.

Parameters:

• conv_module (nn.Module) – The convolutional module to be adjusted.

• reduce (bool (default: None)) – Whether to reduce the dimensionality (e.g., 2D to 1D) or expand it (e.g., 1D to 2D). If None(default), the dimensionality is preserved.

Returns:

A dictionary of adjusted constants.

Return type:

dict

convert(pattern: str, *args, **kwargs) → Module

Convert the dimensionality of the model based on the specified pattern.

Parameters:

pattern (str) – The conversion pattern (e.g., “1d -> 2d”, “2d -> 3d”).

Returns:

The converted PyTorch model.

Return type:

nn.Module

correction_map = {<class 'torch.nn.modules.batchnorm.BatchNorm1d'>: <class 'torch.nn.modules.batchnorm.BatchNorm2d'>, <class 'torch.nn.modules.batchnorm.BatchNorm2d'>: <class 'torch.nn.modules.batchnorm.BatchNorm3d'>, <class 'torch.nn.modules.conv.Conv1d'>: <class 'torch.nn.modules.conv.Conv2d'>, <class 'torch.nn.modules.conv.Conv2d'>: <class 'torch.nn.modules.conv.Conv3d'>, <class 'torch.nn.modules.conv.ConvTranspose1d'>: <class 'torch.nn.modules.conv.ConvTranspose2d'>, <class 'torch.nn.modules.conv.ConvTranspose2d'>: <class 'torch.nn.modules.conv.ConvTranspose3d'>, <class 'torch.nn.modules.conv.LazyConvTranspose1d'>: <class 'torch.nn.modules.conv.LazyConvTranspose2d'>, <class 'torch.nn.modules.conv.LazyConvTranspose2d'>: <class 'torch.nn.modules.conv.LazyConvTranspose3d'>, <class 'torch.nn.modules.pooling.AdaptiveAvgPool1d'>: <class 'torch.nn.modules.pooling.AdaptiveAvgPool2d'>, <class 'torch.nn.modules.pooling.AdaptiveAvgPool2d'>: <class 'torch.nn.modules.pooling.AdaptiveAvgPool3d'>, <class 'torch.nn.modules.pooling.AdaptiveMaxPool1d'>: <class 'torch.nn.modules.pooling.AdaptiveMaxPool2d'>, <class 'torch.nn.modules.pooling.AdaptiveMaxPool2d'>: <class 'torch.nn.modules.pooling.AdaptiveMaxPool3d'>, <class 'torch.nn.modules.pooling.MaxPool1d'>: <class 'torch.nn.modules.pooling.MaxPool2d'>, <class 'torch.nn.modules.pooling.MaxPool2d'>: <class 'torch.nn.modules.pooling.MaxPool3d'>}

dim_correction(reduce: bool)

Correct the dimensionality of the layers identified for conversion.

Parameters:

reduce (bool) – Whether to reduce the dimensionality (e.g., 2D to 1D) or expand it (e.g., 1D to 2D).

find_convertable(model: Module) → tuple[dict, dict]

Find dimension-convertable layers in the given model.

Parameters:

model (nn.Module) – The PyTorch model to be analyzed.

Returns:

A tuple containing two dictionaries: layer_dict and param_dict. layer_dict maps layer names to modules. param_dict maps parameter names to tensors.

Return type:

tuple[dict, dict]

Module contents