Customised related questions

[EEEEE-Q]

Is it possible to build our own model and loss function in this package?

[selimfirat]

Hi @EEEEE-Q ,

Thank you for your question.

Yes, definitely it is possible. Currently, all loss functions and models are compatible with nn.Module. Therefore, you can simply follow Pytorch convention and create model/loss extending nn.Module, and use them in combination with Kaira methods/models.

For future compatibility, you can also extend models.BaseModel and losses.BaseLoss, which also extends nn.Module.

Basic Structure

from kaira.models.base import BaseModel
import torch
import torch.nn as nn

class MyCustomModel(BaseModel):
    """Custom model example."""
    
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, output_dim)
        )
    
    def forward(self, x):
        return self.layers(x)

Creating Custom Models

1. Simple Feedforward Model

from kaira.models.base import BaseModel
from kaira.models.registry import ModelRegistry
import torch
import torch.nn as nn

@ModelRegistry.register_model("my_feedforward")
class FeedforwardModel(BaseModel):
    """Custom feedforward neural network for signal processing.
    
    Parameters
    ----------
    input_dim : int
        Input feature dimension
    hidden_dims : list
        List of hidden layer dimensions
    output_dim : int
        Output dimension
    activation : str, optional
        Activation function ('relu', 'tanh', 'sigmoid')
    dropout : float, optional
        Dropout probability for regularization
    """
    
    def __init__(self, input_dim, hidden_dims, output_dim, 
                 activation='relu', dropout=0.0):
        super().__init__()
        
        self.input_dim = input_dim
        self.output_dim = output_dim
        
        # Build layers
        layers = []
        prev_dim = input_dim
        
        for hidden_dim in hidden_dims:
            layers.extend([
                nn.Linear(prev_dim, hidden_dim),
                self._get_activation(activation),
                nn.Dropout(dropout) if dropout > 0 else nn.Identity()
            ])
            prev_dim = hidden_dim
        
        # Output layer
        layers.append(nn.Linear(prev_dim, output_dim))
        
        self.network = nn.Sequential(*layers)
    
    def _get_activation(self, activation):
        """Get activation function by name."""
        activations = {
            'relu': nn.ReLU(),
            'tanh': nn.Tanh(),
            'sigmoid': nn.Sigmoid(),
            'leaky_relu': nn.LeakyReLU()
        }
        return activations.get(activation, nn.ReLU())
    
    def forward(self, x):
        """Forward pass through the network.
        
        Parameters
        ----------
        x : torch.Tensor
            Input tensor of shape (batch_size, input_dim)
            
        Returns
        -------
        torch.Tensor
            Output tensor of shape (batch_size, output_dim)
        """
        return self.network(x)

2. Convolutional Model for Image Processing

@ModelRegistry.register_model("my_conv_model")
class ConvolutionalModel(BaseModel):
    """Custom CNN for image-based communication tasks.
    
    Parameters
    ----------
    input_channels : int
        Number of input channels
    num_filters : list
        Number of filters for each conv layer
    kernel_sizes : list
        Kernel sizes for each conv layer
    output_dim : int
        Final output dimension
    """
    
    def __init__(self, input_channels, num_filters, kernel_sizes, output_dim):
        super().__init__()
        
        # Convolutional layers
        conv_layers = []
        in_channels = input_channels
        
        for filters, kernel_size in zip(num_filters, kernel_sizes):
            conv_layers.extend([
                nn.Conv2d(in_channels, filters, kernel_size, padding=kernel_size//2),
                nn.BatchNorm2d(filters),
                nn.ReLU(),
                nn.MaxPool2d(2)
            ])
            in_channels = filters
        
        self.conv_layers = nn.Sequential(*conv_layers)
        
        # Global average pooling
        self.global_pool = nn.AdaptiveAvgPool2d(1)
        
        # Classification head
        self.classifier = nn.Linear(in_channels, output_dim)
    
    def forward(self, x):
        """Forward pass through the CNN.
        
        Parameters
        ----------
        x : torch.Tensor
            Input tensor of shape (batch_size, channels, height, width)
            
        Returns
        -------
        torch.Tensor
            Output tensor of shape (batch_size, output_dim)
        """
        # Convolutional features
        features = self.conv_layers(x)
        
        # Global pooling
        pooled = self.global_pool(features)
        pooled = pooled.view(pooled.size(0), -1)
        
        # Classification
        output = self.classifier(pooled)
        
        return output

ChannelAwareBaseModel for CSI

For models that need Channel State Information (CSI), use ChannelAwareBaseModel:

from kaira.models.base import ChannelAwareBaseModel
import torch
import torch.nn as nn

@ModelRegistry.register_model("csi_aware_model")
class CSIAwareModel(ChannelAwareBaseModel):
    """Model that adapts based on channel state information.
    
    Parameters
    ----------
    input_dim : int
        Input signal dimension
    csi_dim : int
        CSI feature dimension
    hidden_dim : int
        Hidden layer dimension
    output_dim : int
        Output dimension
    """
    
    def __init__(self, input_dim, csi_dim, hidden_dim, output_dim):
        super().__init__()
        
        # Signal processing branch
        self.signal_processor = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim)
        )
        
        # CSI processing branch
        self.csi_processor = nn.Sequential(
            nn.Linear(csi_dim, hidden_dim//2),
            nn.ReLU(),
            nn.Linear(hidden_dim//2, hidden_dim)
        )
        
        # Fusion layer
        self.fusion = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, output_dim)
        )
    
    def forward(self, x, csi, **kwargs):
        """Forward pass with CSI adaptation.
        
        Parameters
        ----------
        x : torch.Tensor
            Input signal tensor
        csi : torch.Tensor
            Channel state information tensor
        **kwargs : dict
            Additional keyword arguments
            
        Returns
        -------
        torch.Tensor
            CSI-adapted output
        """
        # Validate and normalize CSI
        csi = self.validate_csi(csi)
        csi = self.normalize_csi(csi, method='minmax')
        
        # Process signal and CSI separately
        signal_features = self.signal_processor(x)
        csi_features = self.csi_processor(csi)
        
        # Fuse features
        combined = torch.cat([signal_features, csi_features], dim=-1)
        output = self.fusion(combined)
        
        return output

BaseLoss Overview

The BaseLoss class provides the foundation for all loss functions in Kaira:

from kaira.losses.base import BaseLoss
import torch
import torch.nn as nn

class BaseLoss(nn.Module, ABC):
    """Base class for all Kaira loss functions."""
    pass

Creating Custom Loss Functions

1. Simple Custom Loss

from kaira.losses.base import BaseLoss
from kaira.losses.registry import LossRegistry
import torch
import torch.nn.functional as F

@LossRegistry.register_loss("my_custom_loss")
class CustomLoss(BaseLoss):
    """Custom loss function for specific use case.
    
    Parameters
    ----------
    alpha : float
        Weighting parameter for first term
    beta : float
        Weighting parameter for second term
    reduction : str
        Reduction method ('mean', 'sum', 'none')
    """
    
    def __init__(self, alpha=1.0, beta=0.5, reduction='mean'):
        super().__init__()
        self.alpha = alpha
        self.beta = beta
        self.reduction = reduction
    
    def forward(self, predictions, targets):
        """Compute the custom loss.
        
        Parameters
        ----------
        predictions : torch.Tensor
            Model predictions
        targets : torch.Tensor
            Ground truth targets
            
        Returns
        -------
        torch.Tensor
            Computed loss value
        """
        # MSE component
        mse_loss = F.mse_loss(predictions, targets, reduction='none')
        
        # L1 component
        l1_loss = F.l1_loss(predictions, targets, reduction='none')
        
        # Combine losses
        combined_loss = self.alpha * mse_loss + self.beta * l1_loss
        
        # Apply reduction
        if self.reduction == 'mean':
            return combined_loss.mean()
        elif self.reduction == 'sum':
            return combined_loss.sum()
        else:
            return combined_loss

2. Perceptual Loss with Feature Matching

@LossRegistry.register_loss("perceptual_loss")
class PerceptualLoss(BaseLoss):
    """Perceptual loss using pretrained feature extractor.
    
    Parameters
    ----------
    feature_extractor : nn.Module
        Pretrained model for feature extraction
    layers : list
        Layer indices to extract features from
    weights : list, optional
        Weights for each layer's contribution
    """
    
    def __init__(self, feature_extractor, layers, weights=None):
        super().__init__()
        
        self.feature_extractor = feature_extractor
        self.layers = layers
        
        if weights is None:
            self.weights = [1.0] * len(layers)
        else:
            assert len(weights) == len(layers)
            self.weights = weights
        
        # Freeze feature extractor
        for param in self.feature_extractor.parameters():
            param.requires_grad = False
        
        self.feature_extractor.eval()
    
    def extract_features(self, x):
        """Extract features from specified layers."""
        features = []
        
        def hook_fn(layer_idx):
            def hook(module, input, output):
                if layer_idx in self.layers:
                    features.append(output)
            return hook
        
        # Register hooks
        hooks = []
        for i, module in enumerate(self.feature_extractor.modules()):
            if i in self.layers:
                hooks.append(module.register_forward_hook(hook_fn(i)))
        
        # Forward pass
        with torch.no_grad():
            _ = self.feature_extractor(x)
        
        # Remove hooks
        for hook in hooks:
            hook.remove()
        
        return features
    
    def forward(self, predictions, targets):
        """Compute perceptual loss.
        
        Parameters
        ----------
        predictions : torch.Tensor
            Model predictions
        targets : torch.Tensor
            Ground truth targets
            
        Returns
        -------
        torch.Tensor
            Perceptual loss value
        """
        # Extract features
        pred_features = self.extract_features(predictions)
        target_features = self.extract_features(targets)
        
        # Compute loss for each layer
        total_loss = 0.0
        for pred_feat, target_feat, weight in zip(pred_features, target_features, self.weights):
            layer_loss = F.mse_loss(pred_feat, target_feat)
            total_loss += weight * layer_loss
        
        return total_loss

3. Adaptive Loss with Uncertainty Weighting

@LossRegistry.register_loss("adaptive_loss")
class AdaptiveLoss(BaseLoss):
    """Adaptive loss that learns to balance multiple objectives.
    
    This loss automatically learns the relative importance of different
    loss components during training.
    """
    
    def __init__(self, loss_functions, initial_weights=None):
        super().__init__()
        
        self.loss_functions = nn.ModuleList(loss_functions)
        
        # Learnable weights (log variance parameters)
        if initial_weights is None:
            initial_weights = [1.0] * len(loss_functions)
        
        self.log_vars = nn.Parameter(
            torch.log(torch.tensor(initial_weights, dtype=torch.float32))
        )
    
    def forward(self, predictions, targets):
        """Compute adaptive weighted loss.
        
        Parameters
        ----------
        predictions : torch.Tensor or list
            Model predictions (can be multiple outputs)
        targets : torch.Tensor or list
            Ground truth targets (can be multiple targets)
            
        Returns
        -------
        torch.Tensor
            Adaptive weighted loss
        """
        total_loss = 0.0
        
        for i, loss_fn in enumerate(self.loss_functions):
            # Compute individual loss
            if isinstance(predictions, (list, tuple)):
                pred = predictions[i] if i < len(predictions) else predictions[0]
            else:
                pred = predictions
                
            if isinstance(targets, (list, tuple)):
                target = targets[i] if i < len(targets) else targets[0]
            else:
                target = targets
            
            loss_i = loss_fn(pred, target)
            
            # Apply uncertainty weighting
            precision = torch.exp(-self.log_vars[i])
            weighted_loss = precision * loss_i + self.log_vars[i]
            
            total_loss += weighted_loss
        
        return total_loss

Model Registration

Register your models for easy discovery and instantiation:

from kaira.models.registry import ModelRegistry

# Register during class definition
@ModelRegistry.register_model("my_model_name")
class MyModel(BaseModel):
    pass

# Or register after definition
ModelRegistry.register_model("alternative_name")(MyModel)

# Create model instances
model = ModelRegistry.create("my_model_name", input_dim=128, output_dim=64)

# List all registered models
available_models = ModelRegistry.list_models()