yamle.models.densenet module

class yamle.models.densenet.EmptyBlock(*args, **kwargs)

Bases: Module

This class defines an empty block that does nothing, but signals where to cache the hidden states.

forward(inputs)

This function applies the forward pass.

Parameters:

inputs (torch.Tensor) – Input tensor.

Return type:

Tensor

training: bool

class yamle.models.densenet.DenseLayer(inplanes, growth_rate, normalization, normalization_kwargs, bn_size, dropout_rate)

Bases: Module

This class defines the primitive of the DenseNet architecture.

Parameters:

• inplanes (int) – Number of input features.

• growth_rate (int) – Number of output features.

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

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

• bn_size (int) – Bottleneck size.

• dropout_rate (float) – Dropout rate.

bn_function(inputs)

This function applies bottleneck function.

Parameters:

inputs (torch.Tensor) – Input tensor.

Return type:

Tensor

forward(inputs)

This function applies the forward pass.

Parameters:

inputs (torch.Tensor) – Input tensor.

Return type:

Tensor

replace_layers_for_quantization()

This function replaces the layers for quantization.

Return type:

None

training: bool

class yamle.models.densenet.Transition(inplanes, outplanes, normalization, normalization_kwargs)

Bases: Module

This class defines the transition layer of the DenseNet architecture.

Parameters:

• inplanes (int) – Number of input features.

• outplanes (int) – Number of output features.

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

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

forward(inputs)

This function applies the forward pass.

Parameters:

inputs (torch.Tensor) – Input tensor.

Return type:

Tensor

training: bool

class yamle.models.densenet.DenseBlock(inplanes, growth_rate, normalization, normalization_kwargs, bn_size, dropout_rate, n_layers)

Bases: Module

This class defines the DenseBlock of the DenseNet architecture.

Parameters:

• inplanes (int) – Number of input features.

• growth_rate (int) – Number of output features.

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

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

• bn_size (int) – Bottleneck size.

• dropout_rate (float) – Dropout rate.

• n_layers (int) – Number of layers.

forward(inputs)

This function applies the forward pass.

Parameters:

inputs (torch.Tensor) – Input tensor.

Return type:

Tensor

training: bool

class yamle.models.densenet.DenseNetModel(layers=[6, 12, 24, 16], depth=4, bn_size=4, growth_rate=32, initial_planes=64, width_multiplier=1.0, normalization='batch', dropout_rate=0.0, *args, **kwargs)

Bases: BaseModel

This class defines the DenseNet architecture as described in the paper. Densely Connected Convolutional Networks: https://arxiv.org/abs/1608.06993

The code is based on the implementation of torchvision: pytorch/vision

Parameters:

• layers (List[int]) – Number of layers in each block.

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

• bn_size (int) – Bottleneck size.

• growth_rate (int) – The growth rate multiplier.

• initial_planes (int) – Number of initial planes.

• width_multiplier (float) – Width multiplier to multiply the initial number of features.

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

• dropout_rate (float) – Dropout rate.

tasks = ['classification', 'regression']

forward(x, staged_output=False, input_kwargs={}, output_kwargs={})

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)

This function is used to get the final layer output.

Return type:

Tensor

add_method_specific_layers(method, **kwargs)

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

Parameters:

method (str) – The method to use.

Return type:

None

training: bool

replace_layers_for_quantization()

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)

Add specific arguments to the parser.

Return type:

ArgumentParser