unet_tapered
¶
U-Net model tapered at the end for low res output.
Module Contents¶
-
class
unet_tapered.
Model
(hparams)¶ Bases:
model.unet.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