yamle.models.vgg module#

class yamle.models.vgg.EmptyBlock(*args, **kwargs)[source]#

Bases: Module

This is just an empty block that does nothing.

forward(x)[source]#

This function does nothing.

Return type:

Tensor

training: bool#
class yamle.models.vgg.VGGBlock(inplanes, planes, pooling, normalization=None, normalization_kwargs={})[source]#

Bases: Module

This class implements a VGG block.

Parameters:

  • inplanes (int) – The number of input planes.

  • planes (int) – The number of output planes.

  • pooling (bool) – Whether to use pooling after the block.

  • normalization (Type[nn.Module]) – The normalization to use.

  • normalization_kwargs (Dict[str, Any]) – The kwargs for the normalization.

forward(x)[source]#

The forward function of the block.

Return type:

Tensor

replace_layers_for_quantization()[source]#

Fuses all the operations in the block.

Return type:

None

training: bool#
class yamle.models.vgg.Blocks(inplanes, planes, blocks, pooling, normalization=None, normalization_kwargs={})[source]#

Bases: Module

A class implementing a stack of VGG blocks.

Parameters:

  • inplanes (int) – The number of input planes.

  • planes (int) – The number of output planes.

  • blocks (int) – The number of blocks.

  • pooling (bool) – Whether to use pooling after the block.

  • normalization (Type[nn.Module]) – The normalization to use.

  • normalization_kwargs (Dict[str, Any]) – The kwargs for the normalization.

forward(x)[source]#

Iterate through the blocks.

Return type:

Tensor

training: bool#
class yamle.models.vgg.VGGModel(layers=[2, 2, 3, 3, 3], depth=None, planes=[64, 128, 256, 512, 512], width_multiplier=1, pooling=[False, True, True, True, True], normalization='batch', *args, **kwargs)[source]#

Bases: BaseModel

This class implements the VGG architecture as described in the paper: Very Deep Convolutional Networks for Large-Scale Image Recognition. and the paper can be found at: https://arxiv.org/abs/1409.1556.

Parameters:

  • layers (List[int]) – The number of layers in the VGG architecture per block.

  • depth (int) – The depth of the network respective to the the length of the layers list.

  • planes (List[int]) – The number of planes in each block of layers.

  • width_multiplier (int) – The width multiplier for the planes list.

  • pooling (List[bool]) – Whether to use pooling after each block of layers.

  • normalization (Optional[str]) – The normalization to use. Can be either ‘batch’, ‘instance’, ‘group’, ‘layer’, or None. Defaults to ‘batch’.

tasks = ['regression', 'classification']#
forward(x, staged_output=False, input_kwargs={}, output_kwargs={})[source]#

The forward function of the ResNet model.

Parameters:

  • x (torch.Tensor) – The input tensor.

  • staged_output (bool) – Whether to return the output of each layer. Defaults to False.

  • input_kwargs (Dict[str, Any]) – The kwargs for the input layer.

  • output_kwargs (Dict[str, Any]) – The kwargs for the output layer.

Return type:

Union[Tensor, Tuple[Tensor, List[Tensor]]]

final_layer(x, **output_kwargs)[source]#

This function is used to get the final layer output.

Return type:

Tensor

add_method_specific_layers(method, **kwargs)[source]#

This method is used to add method specific layers to the model.

Parameters:

method (str) – The method to use.

Return type:

None

replace_layers_for_quantization()[source]#

Fuses all the operations in the network.

In this function we only need to fuse layers that are not in the blocks. e.g. the reshaping layers added by the method.

Return type:

None

static add_specific_args(parent_parser)[source]#

Add specific arguments to the parser.

Return type:

ArgumentParser

training: bool#