echofilter.nn package#

Neural network building blocks.

Subpackages#

Submodules#

echofilter.nn.unet module#

U-Net model.

class echofilter.nn.unet.Down(mode='max', compress_dims=True)[source]#

Bases: torch.nn.modules.module.Module

Downscaling layer, downsampling by a factor of two in one or more dimensions.

forward(x)[source]#

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool#
class echofilter.nn.unet.UNet(in_channels, out_channels, initial_channels=32, bottleneck_channels=None, n_block=4, unet_expansion_factor=2, expand_only_on_down=False, blocks_per_downsample=1, blocks_before_first_downsample=1, always_include_skip_connection=True, deepest_inner='identity', intrablock_expansion=6, se_reduction=4, downsampling_modes='max', upsampling_modes='bilinear', depthwise_separable_conv=True, residual=True, actfn='InplaceReLU', kernel_size=5)[source]#

Bases: torch.nn.modules.module.Module

UNet model.

Parameters

  • in_channels (int) – Number of input channels.

  • out_channels (int) – Number of output channels.

  • initial_channels (int, optional) – Number of latent channels to output from the initial convolution facing the input layer. Default is 32.

  • bottleneck_channels (int, optional) – Number of channels to output from the first block, before the first unet downsampling step can occur. Default is the same as initial_channels.

  • n_block (int, optional) – Number of blocks, both up and down. Default is 4.

  • unet_expansion_factor (int or float, optional) – Channel expansion factor between unet blocks. Default is 2.

  • expand_only_on_down (bool, optional) – Whether to only apply unet_expansion_factor on unet blocks which actually containg a down/up sampling component, and not on vanilla blocks. Default is False.

  • blocks_per_downsample (int or sequence, optional) – Block interval between dowsampling steps in the unet. If this is a sequence, it corresponds to the number of blocks for each spatial dimension. Default is 1.

  • blocks_before_first_downsample (int, optional) – Number of blocks to use before and after the main unet structure. Must be at least 1. Default is 1.

  • always_include_skip_connection (bool, optional) – If True, a skip connection is included between all blocks equally far from the start and end of the UNet. If False, skip connections are only used between downsampling and upsampling operations. Default is True.

  • deepest_inner ({callable, "horizontal_block", "identity", None}, optional) – A layer which should be applied at the deepest part of the network, before the first upsampling step. The parameter should either be a pre-instantiated layer, or the string "horizontal_block", to indicate an additional block as generated by the horizontal_block_factory. If it is the string "identity" or None (default), no additional layer is included at the deepest point before upsampling begins.

  • intrablock_expansion (int or float, optional) – Channel expansion factor within inverse residual block. Default is 6.

  • se_reduction (int or float, optional) – Channel reduction factor within squeeze and excite block. Default is 4.

  • downsampling_modes ({"max", "avg", "stride"} or sequence, optional) – The downsampling mode to use. If this is a string, the same downsampling mode is used for every downsampling step. If it is a sequence, it should contain a string for each downsampling step. If the input sequence is too short, the final value will be used for all remaining downsampling steps. Default is "max".

  • upsampling_modes (str or sequence, optional) – The upsampling mode to use. If this is a string, it must be "conv", or something supported by torch.nn.Upsample; the same upsampling mode is used for every upsampling step. If it is a sequence, it should contain a string for each upsampling step. If the input sequence is too short, the final value will be used for all remaining upsampling steps. Default is "bilinear".

  • depthwise_separable_conv (bool, optional) – Whether to use depthwise separable convolutions in the MBConv block. Otherwise, the depth and pointwise convolutions are fused together into a regular convolution. Default is True.

  • residual (bool, optional) – Whether to use a residual architecture for the MBConv blocks. Default is True.

  • actfn (str, optional) – Name of the activation function to use. Default is "InplaceReLU".

  • kernel_size (int, optional) – Size of convolution kernel to use. Default is 5.

forward(x)[source]#

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool#
class echofilter.nn.unet.UNetBlock(in_channels, horizontal_block_factory, n_block=1, block_expansion_factor=2, expand_only_on_down=False, blocks_per_downsample=1, blocks_before_first_downsample=0, always_include_skip_connection=True, deepest_inner='identity', downsampling_modes='max', upsampling_modes='bilinear', _i_block=0, _i_down=0)[source]#

Bases: torch.nn.modules.module.Module

Create a (cascading set of) UNet block(s).

Each block performs the steps:

  • Store input to be used in skip connection

  • Down step

  • Horizontal block

  • <Recursion>

  • Up step

  • Concatenate with skip connection

  • Horizontal block

Where <Recursion> is a call generating a child UNetBlock instance.

Parameters

  • in_channels (int) – Number of input channels to this block.

  • horizontal_block_factory (callable) – A torch.nn.Module constructor or function which returns a block of layers. The resulting module must accept in_channels and out_channels as its first two arguments.

  • n_block (int, optional) – The number of nested UNetBlocks to use. Default is 1 (no nesting).

  • block_expansion_factor (int or float, optional) – Expansion factor for the number of channels between nested UNetBlocks. Default is 2.

  • expand_only_on_down (bool, optional) – Whether to exand the number of channels only when one of the spatial dimensions is compressed. Default is False.

  • blocks_per_downsample (int or sequence, optional) – How many blocks to include between each downsample operation. This can be a tuple of values for each spatial dimension, or an int which uses the same value for each spatial dimension. Default is 1.

  • blocks_before_first_downsample (int or sequence, optional) – How many blocks to include before the first spatial downsampling occurs. Default is 1.

  • always_include_skip_connection (bool, optional) – If True, a skip connection is included even if no dimensions were downsampled in this block. Default is True.

  • deepest_inner ({callable, "horizontal_block", "identity", None}, optional) – A layer which should be applied at the deepest part of the network, before the first upsampling step. The parameter should either be a pre-instantiated layer, or the string "horizontal_block", to indicate an additional block as generated by the horizontal_block_factory. If it is the string "identity" or None (default), no additional layer is included at the deepest point before upsampling begins.

  • downsampling_modes ({"max", "avg", "stride"} or sequence, optional) – The downsampling mode to use. If this is a string, the same downsampling mode is used for every downsampling step. If it is a sequence, it should contain a string for each downsampling step. If the input sequence is too short, the final value will be used for all remaining downsampling steps. Default is "max".

  • upsampling_modes (str or sequence, optional) – The upsampling mode to use. If this is a string, it must be "conv", or something supported by torch.nn.Upsample; the same upsampling mode is used for every upsampling step. If it is a sequence, it should contain a string for each upsampling step. If the input sequence is too short, the final value will be used for all remaining upsampling steps. Default is "bilinear".

  • _i_block (int, optional) – The current block number. Used internally to track recursion. Default is 0.

  • _i_down (int, optional) – Used internally to track downsampling depth. Default is 0.

Notes

This class is defined recursively, and will instantiate itself as its own child until the number of blocks has been satisfied.

forward(input)[source]#

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool#
class echofilter.nn.unet.Up(in_channels=None, up_dims=True, mode='bilinear')[source]#

Bases: torch.nn.modules.module.Module

Upscaling layer, upsampling by a factor of two in one or more dimensions.

forward(x)[source]#

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool#

echofilter.nn.utils module#

echofilter.nn utility functions.

class echofilter.nn.utils.TensorDict(tensors=None)[source]#

Bases: torch.nn.modules.container.ParameterDict

Hold tensors in a dictionary.

TensorDict can be indexed like a regular Python dictionary, but implements methods such as to which operate on all elements within it.

TensorDict is an ordered dictionary that respects

  • the order of insertion, and

  • in update(), the order of the merged OrderedDict or another TensorDict (the argument to update()).

Note that update() with other unordered mapping types (e.g., Python’s plain dict) does not preserve the order of the merged mapping.

Parameters

tensors (iterable) – A mapping (dictionary) of (string : torch.Tensor) or an iterable of key-value pairs of type (string, torch.Tensor)

detach()[source]#
detach_()[source]#
extra_repr()[source]#

Set the extra representation of the module

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

echofilter.nn.utils.count_parameters(model, only_trainable=True)[source]#

Count the number of (trainable) parameters within a model and its children.

Parameters

  • model (torch.nn.Model) – The model.

  • only_trainable (bool, default=True) – Whether the count should be restricted to only trainable parameters (ones which require grad), otherwise all parameters are included. Default is True.

Returns

Total number of (trainable) parameters possessed by the model.

Return type

int

echofilter.nn.utils.logavgexp(input, dim, keepdim=False, temperature=None, internal_dtype=torch.float32)[source]#

Take the log-average-exp.

Returns the log of meaned exponentials of each row of the input tensor in the given dimension dim. The computation is numerically stabilized.

If keepdim is True, the output tensor is of the same size as input except in the dimension dim where it is of size 1. Otherwise, dim is squeezed (see torch.squeeze()), resulting in the output tensor having 1 fewer dimension.

Parameters

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

  • dim (int) – The dimension to reduce.

  • keepdim (bool, optional) – Whether the output tensor has dim retained or not. Default is False.

  • temperature (float or None, optional) – A temperature which is applied to the logits. Temperatures must be positive. Temperatures greater than 1 make the result closer to the average of input, whilst temperatures 0<t<1 make the result closer to the maximum of input. If None (default) or 1, no temperature is applied.

  • internal_dtype (torch.dtype, optional) – A data type which the input will be cast as before computing the log-sum-exp step. Default is torch.float32.

Returns

The log-average-exp of input.

Return type

torch.Tensor

echofilter.nn.utils.seed_all(seed=None, only_current_gpu=False, mirror_gpus=False)[source]#

Initialize the RNGs for random, numpy, and both CPU and GPU(s) for torch.

Parameters

  • seed (int, optional)) – Seed value to use for the random number generators. If seed is None (default), seeds are picked at random using the methods built in to each RNG.

  • only_current_gpu (bool, default=False) – Whether to only re-seed the current cuda device, or to seed all of them.

  • mirror_gpus (bool, default=False) – Whether all cuda devices should receive the same seed, or different seeds. If mirror_gpus is False and seed is not None, each device receives a different but deterministically determined seed. Default is False.

Notes

Note that we override the settings for the cudnn backend whenever this function is called. If seed is not None, we set:

torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

in order to ensure experimental results behave deterministically and are repeatible. However, enabling deterministic mode may result in an impact on performance. See link for more details. If seed is None, we return the cudnn backend to its performance-optimised default settings of:

torch.backends.cudnn.deterministic = False
torch.backends.cudnn.benchmark = True

echofilter.nn.wrapper module#

Model wrapper.

class echofilter.nn.wrapper.Echofilter(model, top='boundary', bottom='boundary', mapping=None, reduction_ispassive='logavgexp', reduction_isremoved='logavgexp', conditional=False)[source]#

Bases: torch.nn.modules.module.Module

Echofilter logit mapping wrapper.

Parameters

  • model (torch.nn.Module) – The model backbone, which converts inputs to logits.

  • top (str, optional) – Type of output for top line and surface line. If "mask", the top output corresponds to logits, which are converted into probabilities with sigmoid. If "boundary" (default), the output corresponds to logits for the location of the line, which is converted into a probability mask using softmax and cumsum.

  • bottom (str, optional) – As for top, but for the bottom line. Default is "boundary".

  • mapping (dict or None, optional) – Mapping from logit names to output channels provided by model. If None, a default mapping is used. The mapping is stored as self.mapping.

  • reduction_ispassive (str, default="logavgexp") – Method used to reduce the depths dimension for the "logit_is_passive" output.

  • reduction_isremoved (str , default="logavgexp") – Method used to reduce the depths dimension for the "logit_is_removed" output.

  • conditional (bool, optional) – Whether to build a conditional model as well as an unconditional model. If True, there are additional logits in the call output named "x|downfacing" and "x|upfacing", in addition to "x". For instance, "p_is_above_turbulence|downfacing". Default is False.

aliases = [('top', 'turbulence')]#
forward(x, output_device=None)[source]#

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool#
class echofilter.nn.wrapper.EchofilterLoss(reduction='mean', conditional=False, turbulence_mask=1.0, bottom_mask=1.0, removed_segment=1.0, passive=1.0, patch=1.0, overall=0.0, surface=1.0, auxiliary=1.0, ignore_lines_during_passive=False, ignore_lines_during_removed=True, ignore_surface_during_passive=False, ignore_surface_during_removed=True)[source]#

Bases: torch.nn.modules.loss._Loss

Evaluate loss for an Echofilter model.

Parameters

  • reduction ("mean" or "sum", optional) – The reduction method, which is used to collapse batch and timestamp dimensions. Default is "mean".

  • turbulence_mask (float, optional) – Weighting for turbulence line/mask loss term. Default is 1.0.

  • bottom_mask (float, optional) – Weighting for bottom line/mask loss term. Default is 1.0.

  • removed_segment (float, optional) – Weighting for is_removed loss term. Default is 1.0.

  • passive (float, optional) – Weighting for is_passive loss term. Default is 1.0.

  • patch (float, optional) – Weighting for mask_patch loss term. Default is 1.0.

  • overall (float, optional) – Weighting for overall mask loss term. Default is 0.0.

  • surface (float, optional) – Weighting for surface line/mask loss term. Default is 1.0.

  • auxiliary (float, optional) – Weighting for auxiliary loss terms "turbulence-original", "bottom-original", "mask_patches-original", and "mask_patches-ntob". Default is 1.0.

  • ignore_lines_during_passive (bool, optional) – Whether targets for turbulence and bottom lines should be excluded from the loss during passive data collection. Default is True.

  • ignore_lines_during_removed (bool, optional) – Whether targets for turbulence and bottom lines should be excluded from the loss during entirely removed sections. Default is True.

  • ignore_surface_during_passive (bool, optional) – Whether target for the surface line should be excluded from the loss during passive data collection. Default is False.

  • ignore_surface_during_removed (bool, optional) – Whether target for the surface line should be excluded from the loss during entirely removed sections. Default is True.

forward(input, target)[source]#

Construct loss term.

Parameters
reduction: str#