advsecurenet.models package

advsecurenet.models.CustomModels package
- advsecurenet.models.CustomModels.CustomCifar10Model module
  - CustomCifar10Model
    - CustomCifar10Model.forward()
- advsecurenet.models.CustomModels.CustomMnistModel module
  - CustomMnistModel
    - CustomMnistModel.forward()

advsecurenet.models.base_model module

class advsecurenet.models.base_model.BaseModel

Bases: ABC, Module

Abstract class for models.

num_classes

The number of classes in the dataset.

Type:: int

pretrained

Whether to load the pretrained weights or not.

Type:: bool

target_layer

The name of the layer to be used as the target layer.

Type:: str

add_layer(new_layer: Module, position: int = -1, inplace: bool = True) → Module | None

Inserts a new layer into the model at the specified position.

Parameters:

new_layer (nn.Module) – The new layer to be added.
position (int) – The position at which the new layer should be added. By default, it is added at the end.
inplace (bool) – Whether to add the layer in-place or not. If set to False, a new model is created.

Returns:

The new model if inplace is set to False.

Return type:

Optional[nn.Module]

forward(x: Tensor) → Tensor

Forward pass of the model.

Parameters:: x (torch.Tensor) – The input tensor.
Returns:: The output tensor.
Return type:: torch.Tensor

get_layer(layer_name: str) → Module

Retrieve a specific layer module based on its name.

Examples

>>> model = StandardModel(model_name='resnet18', num_classes=10)
>>> model.get_layer('layer1.0.conv1')

get_layer_names() → List[str]: Return a list of layer names in the model.

abstract load_model() → None: Abstract method to load the model. This method should be implemented in derived classes (e.g., StandardModel, CustomModel).

abstract models(): Return a list of available models.

predict(x: Tensor) → Tuple[Tensor, Tensor]

Predicts the class of the input tensor.

Parameters:

x (torch.Tensor) – The input tensor.

Returns:

The predicted class index.
The probability of the predicted class.

Return type:

Tuple[torch.Tensor, torch.Tensor]

save_model(path: str) → None

Save the model to the specified path.

Parameters:: path (str) – The path to save the model.

set_layer(layer_name: str, new_layer: Module)

Replace a specific layer module based on its name with a new module.

Examples

>>> model = StandardModel(model_name='resnet18', num_classes=10)
>>> model.set_layer('layer1.0.conv1', nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False))

advsecurenet.models.base_model.check_model_loaded(func): Wrapper function to check if the model is loaded before calling the decorated function.

advsecurenet.models.custom_model module

class advsecurenet.models.custom_model.CustomModel(config: CustomModelConfig, **kwargs)

Bases: BaseModel

This class is used to load a custom model. It is a subclass of BaseModel.

static available_weights(model_name: str) → list[str]: Not applicable for custom models.

load_model()

Load the custom model. This method is called by the BaseModel constructor. It dynamically imports the custom model based on its name. That is, it assumes that the model class inside the custom model file has the same name as the file i.e., for ‘CustomMnistModel.py’, there should be a class ‘CustomMnistModel’.

Raises:: ValueError – If the model class is not found in the custom model file.

static models()

Returns a list of available custom models from the CustomModels directory.

Returns:: A list of available custom models.
Return type:: List[str]

modify_model()

advsecurenet.models.external_model module

class advsecurenet.models.external_model.ExternalModel(config: ExternalModelConfig, **kwargs)

Bases: BaseModel

This class is used to load external models that are not provided by the package. These models are loaded from external Python files.

load_model(): Loads the external model from the specified path.

models(): Returns a list of available external models.

advsecurenet.models.model_factory module

class advsecurenet.models.model_factory.ModelFactory

Bases: object

This class is a factory class for creating models. It provides a single interface for creating models. It supports both standard models and custom models.

static add_layer(model: Module, new_layer: Module, position: int = -1) → Module

Inserts a new layer into an existing PyTorch model at the specified position. If the model is not a Sequential model, it will be converted into one.

Parameters:

model (nn.Module) – The original model to which the new layer will be added.
new_layer (nn.Module) – The layer to be inserted into the model.
position (int) – The position at which to insert the new layer. If set to -1, the layer is added at the end. Positions are zero-indexed.

Returns:

A new model with the layer added at the specified position.

Return type:

nn.Module

Raises:

ValueError – If the specified position is out of bounds.

static available_custom_models() → list[str]: Returns a list of all available custom models that are created by the user. These models are stored in the ‘advsecurenet/models/CustomModels’ directory.

static available_models() → list[str]: Returns a list of all available models.

static available_standard_models() → list[str]: Returns a list of all available standard models that are supported by torchvision.

static available_weights(model_name: str) → EnumMeta

Returns a list of available weights for the given model_name.

Parameters:: model_name (str) – The name of the model. You can get the list of available models using StandardModel.models().
Returns:: A EnumMeta object containing the available weights for the given model_name.
Return type:: EnumMeta
Raises:: ValueError – If the model_name is not supported.

Note

You can get the list of available weights for a model using list(StandardModel.available_weights(model_name)). This is only applicable for standard models.

Raises:

ValueError – If the model_name is not supported.
ValueError – If the model_name is a custom model.

Examples

>>> ModelFactory.available_weights("resnet18")
<enum 'ResNet18Weights'>
>>> list(ModelFactory.available_weights("resnet18"))
[ResNet18_Weights.IMAGENET1K_V1]

static create_model(config: CreateModelConfig | None = None, **kwargs) → BaseModel

This function creates a model based on the CreateModelConfig.

Parameters:

config (Optional[CreateModelConfig]) – The configuration for creating the model. If not provided, the model will be created with the passed keyword arguments.
fields (CreateModelConfig contains the following) –
- model_name: str
- num_classes: Optional[int] = 1000
- num_input_channels: Optional[int] = 3
- pretrained: Optional[bool] = True
- weights: Optional[str] = “IMAGENET1K_V1”
- custom_models_path: Optional[str] = “CustomModels”
- model_arch_path: Optional[str] = None
- model_weights_path: Optional[str] = None
- is_external: bool = False
- random_seed: Optional[int] = None
**kwargs – Additional keyword arguments to be passed to the model constructor.

Returns:

The created model.

Return type:

BaseModel

Note

If the model is a custom model, the model_name should be the name of the custom model class. For example, ‘CustomMnistModel’. You can use your external model by setting is_external=True in the CreateModelConfig and providing the model_arch_path and model_weights_path.

static infer_model_type(model_name: str) → ModelType

This function infers the model type based on the model_name.

Parameters:: model_name (str) – The name of the model to be loaded. For example, ‘resnet18’ or ‘CustomMnistModel’.
Returns:: The model type of the model_name.
Return type:: ModelType
Raises:: ValueError – If the model_name is not supported by torchvision or is not a custom model.

advsecurenet.models.standard_model module

class advsecurenet.models.standard_model.StandardModel(config: StandardModelConfig, **kwargs)

Bases: BaseModel

This class is used to load standard models from torchvision.models. It supports loading pretrained models and modifying the model after loading.

Parameters:: config (StandardModelConfig) – The configuration for the standard model.

static available_weights(model_name: str) → EnumMeta

Returns a list of available weights for the given model_name.

Parameters:: model_name (str) – The name of the model. You can get the list of available models using StandardModel.models().
Returns:: A EnumMeta object containing the available weights for the given model_name.
Return type:: EnumMeta
Raises:: ValueError – If the model_name is not supported.

Note

You can get the list of available weights for a model using list(StandardModel.available_weights(model_name)).

Examples

>>> StandardModel.available_weights("resnet18")
<enum 'ResNet18Weights'>
>>> list(StandardModel.available_weights("resnet18"))
[ResNet18_Weights.IMAGENET1K_V1]

load_model() → None

Load the model. This method is called by the BaseModel constructor. It loads the model from torchvision.models based on the model_name attribute and sets the model attribute.

Raises:: ValueError – If the model_name is not supported by torchvision.

static models() → list: Returns a list of available standard models from torchvision.models.

modify_model(): Modifies the model after loading. It updates the number of output classes of the pretrained model to the number of classes in the dataset.