CenterBlock¶

class CenterBlock(in_channels)[source]¶

Center block that applies attention mechanism at the bottleneck.

This block is placed at the center of the U-Net architecture (deepest level) to enhance feature representation using attention mechanisms.

Parameters:

in_channels (int) – Number of input channels

Initialize CenterBlock with attention.

Parameters:

in_channels (int) – Number of input channels.

Methods

forward

Forward pass through center block.

Attributes

training

forward(x)[source]¶

Forward pass through center block.

Parameters:

x (torch.Tensor) – Input tensor.

Returns:

Output tensor with attention applied.

Return type:

torch.Tensor