"""Implementation of the DeepJSCC model from :cite:`bourtsoulatze2019deep`."""
from typing import Any, Optional
import torch
from torch import nn
from kaira.models.registry import ModelRegistry
from ..base import BaseModel
class _ConvWithPReLU(nn.Module):
"""Convolutional layer followed by PReLU activation."""
def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, padding: int = 0, *args: Any, **kwargs: Any):
"""Initialize the _ConvWithPReLU module.
Args:
in_channels: Number of input channels.
out_channels: Number of output channels.
kernel_size: Size of the convolutional kernel.
stride: Stride of the convolution.
padding: Padding added to both sides of the input.
*args: Variable positional arguments passed to the base class.
**kwargs: Variable keyword arguments passed to the base class.
"""
super().__init__(*args, **kwargs)
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding)
self.prelu = nn.PReLU()
nn.init.kaiming_normal_(self.conv.weight, mode="fan_out", nonlinearity="leaky_relu")
def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor:
"""Forward pass through convolutional layer and PReLU activation.
Args:
x: Input tensor
*args: Additional positional arguments (unused).
**kwargs: Additional keyword arguments (unused).
Returns:
Activated output tensor
"""
return self.prelu(self.conv(x))
class _TransConvWithPReLU(nn.Module):
"""Transposed convolutional layer followed by activation."""
def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int, padding: int = 0, output_padding: int = 0, activate: Optional[nn.Module] = None, *args: Any, **kwargs: Any):
"""Initialize the _TransConvWithPReLU module.
Args:
in_channels: Number of input channels.
out_channels: Number of output channels.
kernel_size: Size of the transposed convolutional kernel.
stride: Stride of the transposed convolution.
padding: Padding added to both sides of the input.
output_padding: Additional size added to one side of the output shape.
activate: Activation function to use. Defaults to PReLU.
*args: Variable positional arguments passed to the base class.
**kwargs: Variable keyword arguments passed to the base class.
"""
super().__init__(*args, **kwargs)
self.transconv = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding, output_padding)
# Use PReLU by default if no activation is specified
self.activate = activate if activate is not None else nn.PReLU()
if isinstance(self.activate, nn.PReLU):
nn.init.kaiming_normal_(self.transconv.weight, mode="fan_out", nonlinearity="leaky_relu")
else:
nn.init.xavier_normal_(self.transconv.weight)
def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor:
"""Forward pass through transposed convolutional layer and activation.
Args:
x: Input tensor
*args: Additional positional arguments (unused).
**kwargs: Additional keyword arguments (unused).
Returns:
Activated output tensor
"""
return self.activate(self.transconv(x))
[docs]
@ModelRegistry.register_model()
class Bourtsoulatze2019DeepJSCCEncoder(BaseModel):
"""DeepJSCC encoder model from :cite:`bourtsoulatze2019deep`.
This model encodes the input image into a latent representation for transmission.
Args:
num_transmitted_filters: Number of filters in the final encoding layer
"""
[docs]
def __init__(self, num_transmitted_filters: int, *args: Any, **kwargs: Any):
"""Initialize the Bourtsoulatze2019DeepJSCCEncoder.
Args:
num_transmitted_filters: Number of filters in the final encoding layer.
*args: Variable positional arguments passed to the base class.
**kwargs: Variable keyword arguments passed to the base class.
"""
super().__init__(*args, **kwargs)
self.model = nn.Sequential(
_ConvWithPReLU(in_channels=3, out_channels=16, kernel_size=5, stride=2, padding=2),
_ConvWithPReLU(in_channels=16, out_channels=32, kernel_size=5, stride=2, padding=2),
_ConvWithPReLU(in_channels=32, out_channels=32, kernel_size=5, padding=2),
_ConvWithPReLU(in_channels=32, out_channels=32, kernel_size=5, padding=2),
_ConvWithPReLU(in_channels=32, out_channels=num_transmitted_filters, kernel_size=5, padding=2),
)
[docs]
def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor:
"""Forward pass through the encoder.
Args:
x: Input image tensor of shape (B, 3, H, W)
*args: Additional positional arguments (unused).
**kwargs: Additional keyword arguments (unused).
Returns:
Encoded representation of shape (B, num_transmitted_filters, H//4, W//4)
"""
return self.model(x)
[docs]
@ModelRegistry.register_model()
class Bourtsoulatze2019DeepJSCCDecoder(BaseModel):
"""DeepJSCC decoder model from :cite:`bourtsoulatze2019deep`.
This model decodes the transmitted representation back into an image.
Args:
num_transmitted_filters: Number of filters in the transmitted representation
"""
[docs]
def __init__(self, num_transmitted_filters: int, *args: Any, **kwargs: Any):
"""Initialize the Bourtsoulatze2019DeepJSCCDecoder.
Args:
num_transmitted_filters: Number of filters in the transmitted representation.
*args: Variable positional arguments passed to the base class.
**kwargs: Variable keyword arguments passed to the base class.
"""
super().__init__(*args, **kwargs)
self.model = nn.Sequential(
_TransConvWithPReLU(in_channels=num_transmitted_filters, out_channels=32, kernel_size=5, stride=1, padding=2),
_TransConvWithPReLU(in_channels=32, out_channels=32, kernel_size=5, stride=1, padding=2),
_TransConvWithPReLU(in_channels=32, out_channels=32, kernel_size=5, stride=1, padding=2),
_TransConvWithPReLU(in_channels=32, out_channels=16, kernel_size=5, stride=2, padding=2, output_padding=1),
_TransConvWithPReLU(in_channels=16, out_channels=3, kernel_size=5, stride=2, padding=2, output_padding=1, activate=nn.Sigmoid()),
)
[docs]
def forward(self, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor:
"""Forward pass through the decoder.
Args:
x: Encoded representation of shape (B, num_transmitted_filters, H//4, W//4)
*args: Additional positional arguments (unused).
**kwargs: Additional keyword arguments (unused).
Returns:
Decoded image tensor of shape (B, 3, H, W)
"""
return self.model(x)