kaira.models.components.AFModule

Inheritance diagram for AFModule

class kaira.models.components.AFModule(N, csi_length, *args: Any, **kwargs: Any)

Bases: BaseModel

AFModule: Attention-Feature Module [Xu et al., 2021].

This module implements a an attention mechanism that recalibrates feature maps by explicitly modeling interdependencies between channel state information and the input features. This module allows the same model to be used during training and testing across channels with different signal-to-noise ratio without significant performance degradation.

Methods

__init__	Initialize the AFModule.
forward	Forward pass through the AFModule.

Examples using kaira.models.components.AFModule

Attention-Feature Module (AFModule)

Attention-Feature Module (AFModule)

__init__(N, csi_length, *args: Any, **kwargs: Any)

Initialize the AFModule.

Parameters:

• N (int) – The number of input and output features.

• csi_length (int) – The length of the channel state information.

• *args – Variable positional arguments passed to the base class.

• **kwargs – Variable keyword arguments passed to the base class.

forward(x: Tensor, csi: Tensor, *args: Any, **kwargs: Any) → Tensor

Forward pass through the AFModule.

Parameters:

• x (torch.Tensor) – The input tensor.

• csi (torch.Tensor) – Channel State Information tensor.

• *args – Additional positional arguments (unused).

• **kwargs – Additional keyword arguments (unused).

Returns:

The output tensor after applying the attention mechanism.

Return type:

torch.Tensor