UNet#

class torch_models.Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros', device=None, dtype=None)#

Inherents the 2D convolution layer fron torch.nn.conv2d

Parameters#

in_channelsint

Number of input channels

out_channelsint

Number of channels produced by the convolution

kernel_sizeint

Size of the kernel matrix

strideint, optional

Size of stride - how far the filter moves in each step, by default 1

paddingstr or int, optional

Padding strategy, by default ‘default’

dilationint, optional

Spacing between kernel elements, by default 1

groupsint, optional

Number of blocked connections from input channels to output channels, by default 1

biasbool, optional

Whether to add a learnable bias to the output, by default True

padding_modestr, optional

Padding method, by default ‘zeros’

devicetorch.device, optional

The PyTorch device destination , by default None

dtypendarray.dtype

ndarray dtype, by default None

class torch_models.Conv2dBNReLU(in_channels, out_channels, kernel_size, stride=1, padding='default', dilation=1, groups=1, bias=None, norm_layer='batch', activation=ReLU(inplace=True), dropout_rate=0.0)#

Create2 a Conv2d->BN->ReLU layer.

Parameters#

in_channelsint

Number of input channels

out_channelsint

Number of channels produced by the convolution

kernel_sizeint

Size of the kernel matrix

strideint, optional

Size of stride - how far the filter moves in each step, by default 1

paddingstr, optional

Padding strategy, by default ‘default’

dilationint, optional

Spacing between kernel elements, by default 1

groupsint, optional

Number of blocked connections from input channels to output channels, by default 1

biasbool, optional

Whether to add a learnable bias to the output, by default None

norm_layerstr, optional

Name of the normalization layers, by default ‘batch’

activationnn alike, optional

Activation function, by default nn.ReLU(inplace=True)

dropout_ratefloat, optional

Dropout rate, by default 0.0

Notes#

padding:

‘default’ (default): torch standard symmetric padding with (kernel_size - 1) // 2. int: symmetric padding to pass to nn.Conv2d(padding=padding) “same”: tf padding=”same”, asymmetric for even kernel (l_0, r_1), etc) “valid”: tf padding=”valid”, same as padding=0

For futher reference please consult nn.conv2d documentation: https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html

class torch_models.ConvBNReLU(conv, norm_layer=None, activation=None, dropout_rate=0.0)#

Initializes internal Module state, shared by both nn.Module and ScriptModule.

class torch_models.IoU(mode='iou', axis=1, eps=0.0)#

Return a matrix of [batch * num_classes].

Parameters#

modestr, optional

Evaluation metrics, by default ‘iou’. Options [‘iou’, ‘dice’]

axisint, optional

The direction to calculate IoU (row or column-wise), by default 1

epsfloat, optional

Used to prevent zero division error, by default 0.

Notes#

In order to separate from iou=0, function WILL return NaN if both y_true and y_pred are 0. Need further treatment to remove nan in either loss function or matrix.

forward(y_pred, y_true)#

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.

class torch_models.ResizeConv2d(in_channels, out_channels, stride=2, conv=(<class 'torch_models.Conv2d'>, {'kernel_size': 3, 'padding': 'default'}), mode='bilinear')#

Upsampling layer with resize convolution.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

class torch_models.SoftDiceLoss(weight=None, ignore_index=[], reduction='mean', mode='dice', axis=1, eps=0.0, use_positive=False)#

Return a matrix of [batch * num_classes].

Parameters#

modestr, optional

Evaluation metrics, by default ‘iou’. Options [‘iou’, ‘dice’]

axisint, optional

The direction to calculate IoU (row or column-wise), by default 1

epsfloat, optional

Used to prevent zero division error, by default 0.

Notes#

In order to separate from iou=0, function WILL return NaN if both y_true and y_pred are 0. Need further treatment to remove nan in either loss function or matrix.

forward(y_pred, y_true)#

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.

class torch_models.SubPixelConv2d(in_channels, out_channels, stride=2, conv=(<class 'torch_models.Conv2d'>, {'kernel_size': 3, 'padding': 'default'}))#

Upsampling layer with better modelling power.

Notes#

Sub-pixel convolution usually gives better result than resize convolution. Use incr init (and weight_norm) to avoid checkboard artifact, ref: https://arxiv.org/pdf/1707.02937.pdf May combine with (https://arxiv.org/pdf/1806.02658.pdf): nn.Sequential(nn.LeakyReLU(inplace=True), nn.ReplicationPad2d((1,0,1,0)), nn.AvgPool2d(2, stride=1)) to generate non-checkboard artifact image.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

icnr_(x)#

ICNR init of conv weight.

class torch_models.UNet(num_classes, scale_factor=1, resize_mode='bilinear', encoder={}, decoder={}, out_channels=None, return_logits=False, **kwargs)#

Abstract UNet class.

Parameters#

num_classesint

Number of classes in the output layer

scale_factorint, optional

How much to rescale compared to the original size, by default 1

resize_modestr, optional

By default ‘bilinear’

encodernn.Module or dictionary, optional

Calls UNetfeatures(**encoder)), by default {}

decodernn.Module or dictionary, optional

Calls self.default_decoder(**decoder), by default {}

out_channels_type_, optional

Calls encoder.feature_channels(), by default None

return_logitsbool, optional

Whether to return the logits from the layers, by default False

forward(x)#

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.

get_encoder(**kwargs)#

Initializes class UNetFeatures

Returns#

UNetFeatures

Defining a part of the U-Net model

class torch_models.UNetFeatures(in_channels=3, n_channels=32, n_downsampling=4, pool='maxpool', norm_layer='batch')#

Initializes the U-Net encoder

Parameters#

in_channelsint, optional

Number of input channels, by default 3

n_channelsint, optional

Number of channels produced by the convolution, by default 32

n_downsamplingint, optional

Number of downsampling layers, by default 4

poolstr, optional

Pooling strategy, by default ‘maxpool’

norm_layerstr, optional

Name of the normalization layers, by default ‘batch’

forward(x)#

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.

torch_models.get_norm_layer_and_bias(norm_layer='batch', use_bias=None)#

Return a normalization layer and set up use_bias for convoluation layers.

Parameters#

norm_layerstr or nn alike, optional

If str, the name of the normalization layer: [batch, instance] If nn alike: nn.BatchNorm2d, nn.InstanceNorm2d If None: no batch norm

use_biasbool, optional

Whether to add a learnable bias to the output, by default None

Returns#

norm_layernn alike

One of {nn.BatchNorm2d, nn.InstanceNorm2d}

Raises#

NotImplementedError

Raised when normalization is neither “batch” or “instance”

Notes#

For BatchNorm: use learnable affine parameters. (affine=True)

track running statistics (mean/stddev). (track_running_stats=True) do not use bias in previous convolution layer. (use_bias=False)

For InstanceNorm: do not use learnable affine parameters. (affine=False)

do not track running statistics. (track_running_stats=False) use bias in previous convolution layer. (use_bias=True)

Examples#

get_norm_layer_and_bias(‘batch’, None) -> affine=True, track_running_stats=True, False get_norm_layer_and_bias(‘batch’, True) -> affine=True, track_running_stats=True, True get_norm_layer_and_bias(‘instance’, None) -> affine=False, track_running_stats=False, True get_norm_layer_and_bias(‘instance’, False) -> affine=False, track_running_stats=False, False get_norm_layer_and_bias(None, None) -> None, True get_norm_layer_and_bias(None, False) -> None, False get_norm_layer_and_bias(nn.BatchNorm2d, None) -> BatchNorm2d, False get_norm_layer_and_bias(nn.BatchNorm2d, True) -> BatchNorm2d, True get_norm_layer_and_bias(nn.InstanceNorm2d, None) -> InstanceNorm2d, True get_norm_layer_and_bias(nn.InstanceNorm2d, False) -> InstanceNorm2d, False