|
|
from torch import nn |
|
|
from typing import Tuple |
|
|
import torch |
|
|
|
|
|
|
|
|
class ChannelAdapter(nn.Module): |
|
|
"""Nonlinear encoder for channel condition tokens. |
|
|
|
|
|
Creates embeddings for SNR, delay spread, and Doppler shift parameters. |
|
|
Each embedding is conditioned on a single real value and is the output of an MLP |
|
|
""" |
|
|
|
|
|
def __init__(self, hidden_sizes: Tuple[int, int, int]): |
|
|
"""Initialize the token encoder modules. |
|
|
|
|
|
Args: |
|
|
hidden_sizes: Tuple of hidden layer dimensions (h1, h2, h3) for the MLP encoders |
|
|
""" |
|
|
super().__init__() |
|
|
self.snr_encoder = self._create_mlp(hidden_sizes) |
|
|
self.ds_encoder = self._create_mlp(hidden_sizes) |
|
|
self.dop_encoder = self._create_mlp(hidden_sizes) |
|
|
|
|
|
@staticmethod |
|
|
def _create_mlp(hidden_sizes: Tuple[int, int, int]) -> nn.Sequential: |
|
|
"""Create a multi-layer perceptron with specified dimensions. |
|
|
|
|
|
Args: |
|
|
hidden_sizes: Tuple of hidden layer dimensions (h1, h2, h3) |
|
|
|
|
|
Returns: |
|
|
Sequential MLP model with ReLU activations between linear layers |
|
|
""" |
|
|
return nn.Sequential( |
|
|
nn.Linear(1, hidden_sizes[0]), |
|
|
nn.ReLU(), |
|
|
nn.Linear(hidden_sizes[0], hidden_sizes[1]), |
|
|
nn.ReLU(), |
|
|
nn.Linear(hidden_sizes[1], hidden_sizes[2]) |
|
|
) |
|
|
|
|
|
def forward( |
|
|
self, |
|
|
snr: torch.Tensor, |
|
|
delay_spread: torch.Tensor, |
|
|
doppler_shift: torch.Tensor |
|
|
) -> torch.Tensor: |
|
|
"""Create token embeddings from channel conditions. |
|
|
|
|
|
Args: |
|
|
snr: Signal-to-Noise Ratio tensor of shape (batch_size, 1) |
|
|
delay_spread: Delay spread tensor of shape (batch_size, 1) |
|
|
doppler_shift: Doppler shift tensor of shape (batch_size, 1) |
|
|
|
|
|
Returns: |
|
|
Concatenated token embeddings of shape (batch_size, 3, 6) |
|
|
""" |
|
|
batch_size = snr.shape[0] |
|
|
snr_emb = torch.reshape(self.snr_encoder(snr), (batch_size, -1, 2)) |
|
|
ds_emb = torch.reshape(self.ds_encoder(delay_spread), (batch_size, -1, 2)) |
|
|
dop_emb = torch.reshape(self.dop_encoder(doppler_shift), (batch_size, -1, 2)) |
|
|
return torch.cat((snr_emb, ds_emb, dop_emb), dim=2) |
|
|
|
