unet_tapered

U-Net model tapered at the end for low res output.

Module Contents

class unet_tapered.Model(hparams)

Bases: model.unet.Model

Inheritance diagram of unet_tapered.Model

This class implements modified U-Net module by removing the up-sampling layers once activation resolution becomes 1/4th the resolution of input. After the removal, feature compression layers are added which keep the resolution constant all along. The layers in the tapered end additionally have skip connections similar to DenseNet. It is equivalent to PyTorch’s nn.Module in all aspects.

Parameters

LightningModule (nn.Module) – The Pytorch-Lightning module derived from nn.module with useful hooks

forward(self, x)

Does the forward pass on the model.

Parameters

x (torch.Tensor) – Input tensor batch.

Returns

Output activations.

Return type

torch.Tensor