kaira.models.image.Yilmaz2024DeepJSCCWZConditionalEncoder

Inheritance diagram for Yilmaz2024DeepJSCCWZConditionalEncoder

class kaira.models.image.Yilmaz2024DeepJSCCWZConditionalEncoder(N: int, M: int, *args: Any, **kwargs: Any)

Bases: ChannelAwareBaseModel

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

This variant of the DeepJSCC-WZ encoder actively incorporates side information during the encoding process. This model is designed for scenarios where side information is available at both encoder and decoder, serving as an upper bound for performance comparison.

The conditional encoder features three processing paths: - g_a: Main encoding path that fuses the input with side information features - g_a2: Processing path for side information for the decoder - g_a3: Auxiliary path for feature extraction from side information for the encoder

By leveraging correlations between the main signal and side information at encoding time, this model achieves more efficient compression and better reconstruction quality compared to the standard DeepJSCC-WZ model, at the cost of requiring side information during encoding.

Architecture highlights: - Early fusion of input and side information (6-channel input) - Multi-scale feature fusion with side information - 4 stages of downsampling (16× spatial reduction) - Channel-adaptive processing with AFModule

Methods

__init__	Initialize the conditional DeepJSCC-WZ encoder.
create_csi_for_submodules	Create appropriate CSI tensors for multiple submodules.
extract_csi_features	Extract common features from CSI tensor for analysis.
extract_csi_from_channel_output	Extract CSI from channel output if available.
format_csi_for_channel	Format CSI tensor for passing to channels that expect specific formats.
forward	Encode the input image with conditional side information.
forward_csi_to_submodules	Helper method to consistently pass CSI to submodules.
normalize_csi	Normalize CSI values to a specified range.
transform_csi	Transform CSI tensor to match target shape requirements.
validate_csi	Validate and ensure CSI tensor is in the correct format.

__init__(N: int, M: int, *args: Any, **kwargs: Any) → None

Initialize the conditional DeepJSCC-WZ 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, x_side: Tensor, csi: Tensor, *args: Any, **kwargs: Any) → Tensor

Encode the input image with conditional side information.

This method processes both the main input image and side information in parallel, fusing features from the side information stream into the main encoding path at multiple scales. The side information is available at the encoder, allowing for more efficient compression compared to the standard DeepJSCC-WZ model.

Parameters:

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

• x_side (torch.Tensor) – Side information tensor of shape [B, 3, H, W] used during encoding.

• csi (torch.Tensor) – Channel state information tensor of shape [B, 1, 1, 1].

• *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

create_csi_for_submodules(csi: Tensor, num_modules: int) → List[Tensor]

Create appropriate CSI tensors for multiple submodules.

Parameters:

• csi (torch.Tensor) – Original CSI tensor

• num_modules (int) – Number of submodules that need CSI

Returns:

List of CSI tensors for each submodule

Return type:

List[torch.Tensor]

extract_csi_features(csi: Tensor) → Dict[str, Tensor]

Extract common features from CSI tensor for analysis.

Parameters:

csi (torch.Tensor) – The CSI tensor to analyze

Returns:

Dictionary containing extracted features

Return type:

Dict[str, torch.Tensor]

static extract_csi_from_channel_output(channel_output: Any) → Tensor | None

Extract CSI from channel output if available.

Some channels return both the transmitted signal and CSI information. This static method provides a standardized way to extract CSI from various channel output formats.

Parameters:

channel_output – Output from a channel, which may contain CSI

Returns:

Extracted CSI tensor if available, None otherwise

Return type:

Optional[torch.Tensor]

static format_csi_for_channel(csi: Tensor, channel_format: str = 'tensor') → Any

Format CSI tensor for passing to channels that expect specific formats.

Parameters:

• csi (torch.Tensor) – CSI tensor to format

• channel_format (str) – Expected format (“tensor”, “dict”, “kwargs”)

Returns:

Formatted CSI in the requested format

Return type:

Any

forward_csi_to_submodules(csi: Tensor, modules: List[BaseModel], *args, **kwargs) → List[Any]

Helper method to consistently pass CSI to submodules.

This method facilitates passing CSI to multiple submodules that require channel state information, ensuring consistent handling across the model.

Parameters:

• csi (torch.Tensor) – Channel state information tensor

• modules (List[BaseModel]) – List of modules to apply

• *args – Positional arguments to pass to modules

• **kwargs – Keyword arguments to pass to modules

Returns:

List of outputs from each module

Return type:

List[Any]

normalize_csi(csi: Tensor, method: str = 'minmax', target_range: tuple = (0.0, 1.0)) → Tensor

Normalize CSI values to a specified range.

Parameters:

• csi (torch.Tensor) – The CSI tensor to normalize

• method (str) – Normalization method. Options: “minmax”, “zscore”, “none”

• target_range (tuple) – Target range for minmax normalization (min, max)

Returns:

Normalized CSI tensor

Return type:

torch.Tensor

Raises:

ValueError – If normalization method is not supported

transform_csi(csi: Tensor, target_shape: Size) → Tensor

Transform CSI tensor to match target shape requirements.

Parameters:

• csi (torch.Tensor) – The CSI tensor to transform

• target_shape (torch.Size) – Target shape for the CSI tensor

Returns:

Transformed CSI tensor

Return type:

torch.Tensor

validate_csi(csi: Tensor, expected_shape: Size | None = None) → Tensor

Validate and ensure CSI tensor is in the correct format.

Parameters:

• csi (torch.Tensor) – The CSI tensor to validate

• expected_shape (Optional[torch.Size]) – Expected shape for the CSI tensor. If None, uses cached shape or infers from tensor.

Returns:

Validated CSI tensor

Return type:

torch.Tensor

Raises:

• ValueError – If CSI tensor is invalid or has incorrect shape

• TypeError – If CSI is not a tensor