kaira.models.image.Yilmaz2024DeepJSCCWZSmallEncoder

Inheritance diagram of Yilmaz2024DeepJSCCWZSmallEncoder

Inheritance diagram for Yilmaz2024DeepJSCCWZSmallEncoder

class kaira.models.image.Yilmaz2024DeepJSCCWZSmallEncoder(N: int, M: int, *args: Any, **kwargs: Any)[source]

Bases: BaseModel

DeepJSCC-WZ-sm Encoder Module [Yilmaz et al., 2024].

This is a lightweight version of the DeepJSCC-WZ encoder that transforms input images into a compressed latent representation suitable for transmission over noisy channels. The encoder consists of a series of residual blocks with downsampling, attention modules, and adaptive feature modules that incorporate channel state information (CSI).

DeepJSCC-WZ-sm shares encoder parameters for encoding image at the transmitter and encoding side information at the receiver, resulting in a parameter-efficient design while maintaining competitive performance.

Architecture highlights: - 4 stages of downsampling (factor of 16 total spatial reduction) - Attention mechanisms to capture important features - AFModule layers that adapt features based on channel conditions - Progressive compression: 3×H×W → M×(H/16)×(W/16) - Channel-aware design through CSI conditioning

Methods

__init__

Initialize the DeepJSCC-WZ-sm encoder.

forward

Process input image through the encoder.
__init__(N: int, M: int, *args: Any, **kwargs: Any) → None[source]

Initialize the DeepJSCC-WZ-sm encoder.

Parameters:

  • N (int) – Number of intermediate channels in the residual blocks. Controls the network capacity and feature dimension.

  • M (int) – Number of output channels in the final latent representation. Determines the compression rate and bandwidth usage.

  • *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[source]

Process input image through the encoder.

Parameters:

  • x (torch.Tensor) – Input image tensor of shape [B, 3, H, W].

  • csi (torch.Tensor) – Channel state information tensor of shape [B, 1, 1, 1]. Contains SNR or other channel quality indicators.

  • *args – Additional positional arguments (passed to internal layers).

  • **kwargs – Additional keyword arguments (passed to internal layers).

Returns:

Encoded representation ready for transmission.

Shape: [B, M, H/16, W/16], where M is the number of channels specified during initialization.

Return type:

torch.Tensor