yamle.models.unet module#

class yamle.models.unet.DownBlock(in_channels, out_channels, normalization)[source]#

Bases: Module

This class is used to create a down block of the UNet model.

Parameters:

  • in_channels (int) – The number of input channels.

  • out_channels (int) – The number of output channels.

  • normalization (str) – The normalization to use.

forward(x)[source]#

The forward function of the down block.

Return type:

Tensor

training: bool#
class yamle.models.unet.UpBlock(in_channels, out_channels, normalization)[source]#

Bases: Module

This class is used to create an up block of the UNet model.

Parameters:

  • in_channels (int) – The number of input channels.

  • out_channels (int) – The number of output channels.

  • normalization (str) – The normalization to use.

forward(x, skip)[source]#

The forward function of the up block.

Return type:

Tensor

training: bool#
class yamle.models.unet.UNetModel(init_features=32, normalization='batch', *args, **kwargs)[source]#

Bases: BaseModel

This class is used to create the UNet model.

Parameters:

  • init_features (int) – The number of initial features.

  • normalization (str) – The type of normalization to use. Defaults to batch. Choices are batch, layer, instance or None.

tasks = ['depth_estimation', 'reconstruction', 'segmentation']#
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

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

The forward function of the UNet 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 input arguments to pass to the input layer. Defaults to {}.

  • output_kwargs (Dict[str, Any]) – The input arguments to pass to the output layer. Defaults to {}.

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

static add_specific_args(parent_parser)[source]#

This function is used to add specific arguments to the parser.

Return type:

ArgumentParser

training: bool#