kaira.models.components.ConvEncoder

Inheritance diagram for ConvEncoder

class kaira.models.components.ConvEncoder(in_channels: int, out_features: int, hidden_dims: List[int] | None = None, kernel_size: int = 3, stride: int = 2, padding: int = 1, activation: Module | None = None, *args: Any, **kwargs: Any)

Bases: BaseModel

Convolutional Neural Network (CNN) Encoder for image transmission systems.

This module implements a CNN-based encoder that maps input images to encoded signals suitable for transmission over a communication channel.

Methods

__init__	Initialize the ConvEncoder.
forward	Forward pass of the ConvEncoder.

__init__(in_channels: int, out_features: int, hidden_dims: List[int] | None = None, kernel_size: int = 3, stride: int = 2, padding: int = 1, activation: Module | None = None, *args: Any, **kwargs: Any)

Initialize the ConvEncoder.

Parameters:

• in_channels (int) – Number of input channels in the image.

• out_features (int) – Dimensionality of the output encoded signals.

• hidden_dims (List[int], optional) – List of feature dimensions for hidden layers. If None, default dimensions [16, 32, 64] will be used.

• kernel_size (int, optional) – Kernel size for convolutions. Default is 3.

• stride (int, optional) – Stride for convolutions. Default is 2.

• padding (int, optional) – Padding for convolutions. Default is 1.

• activation (nn.Module, optional) – Activation function to use. If None, ReLU is used.

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

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

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

Forward pass of the ConvEncoder.

Parameters:

• x (torch.Tensor) – Input image tensor of shape (batch_size, in_channels, height, width).

• *args – Additional positional arguments (unused).

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

Returns:

Output tensor of shape (batch_size, out_features).

Return type:

torch.Tensor