dtc_super_resolution.rcm_emulator.ml_module.basic_unet.unet_parts

Parts of the U-Net model.

Classes

DoubleConv(in_channels, out_channels[, ...])	(convolution => [BN] => ReLU) * 2.
Down(in_channels, out_channels)	Downscaling with double conv then maxpool.
InitialConv(in_channels, out_channels, ...)	Initial convolution to adapt spatial dimensions.
LinearUpsamplingZ([in_channels, out_channels])	Linear upsampling of z to match the number of channels of downsampled x.
OutConv(in_channels, out_channels)	Final convolution to reduce channels to 1.
Up(in_channels, out_channels[, bilinear])	Upscaling then double conv.
UpLastConv(in_channels, out_channels, ...[, ...])	Upscaling then double conv without concatenating with skip connection.

class DoubleConv(in_channels: int, out_channels: int, mid_channels: int = 0)

Bases: Module

(convolution => [BN] => ReLU) * 2.

forward(x: Tensor) → Tensor

Forward pass for the DoubleConv module.

class Down(in_channels: int, out_channels: int)

Bases: Module

Downscaling with double conv then maxpool.

forward(x: Tensor) → tuple[Tensor, Tensor]

Forward pass for the Down module.

class InitialConv(in_channels: int, out_channels: int, size_in: int, size_out: int)

Bases: Module

Initial convolution to adapt spatial dimensions.

forward(x: Tensor) → Tensor

Forward pass for the InitialConv module.

class LinearUpsamplingZ(in_channels: int = 16, out_channels: int = 1024)

Bases: Module

Linear upsampling of z to match the number of channels of downsampled x.

forward(z: Tensor) → Tensor

Forward pass for linear upsampling of z.

class OutConv(in_channels: int, out_channels: int)

Bases: Module

Final convolution to reduce channels to 1.

forward(x: Tensor) → Tensor

Forward pass for the OutConv module.

class Up(in_channels: int, out_channels: int, bilinear: bool = True)

Bases: Module

Upscaling then double conv.

forward(x1: Tensor, x2: Tensor) → Tensor

Forward pass for the Up module.

class UpLastConv(in_channels: int, out_channels: int, mid_channels: int, bilinear: bool = True)

Bases: Module

Upscaling then double conv without concatenating with skip connection.

forward(x: Tensor, zsize: int = 64) → Tensor

Forward pass for the UpLastConv module.