Output networks updated

Author: felix-e-h-p

pvnet/models/multimodal/linear_networks/output_networks.py

@@ -11,13 +11,18 @@
 
 import torch
 import torch.nn.functional as F
+import logging
 from torch import nn
 from abc import ABC, abstractmethod
 from typing import Optional, List, Dict, Union
 
 from pvnet.models.multimodal.linear_networks.basic_blocks import AbstractLinearNetwork
 
 
+logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger('output_networks')
+
+
 class DynamicOutputNetwork(AbstractLinearNetwork):
     """ Dynamic output network definition """
 
@@ -37,14 +42,20 @@ def __init__(
     ):
         # Initialisation of dynamic output network
         super().__init__(in_features=in_features, out_features=out_features)
+        logger.info(f"Initialising DynamicOutputNetwork with in_features={in_features}, out_features={out_features}")
+        logger.debug(f"Configuration - dropout: {dropout}, layer_norm: {use_layer_norm}, residual: {use_residual}")
+
         self.out_features = out_features
 
         # Default hidden architecture
         # h_i ∈ ℝ^{d_i}, where d_i = [2n, n]
         if hidden_dims is None:
             hidden_dims = [in_features * 2, in_features]
-            
+            logger.debug(f"Using default hidden dimensions: {hidden_dims}")
+
         if any(dim <= 0 for dim in hidden_dims):
+            error_msg = f"Invalid hidden dimensions: {hidden_dims}"
+            logger.error(error_msg)
             raise ValueError("hidden_dims must be positive")
 
         # Construction of network layers - config
@@ -57,19 +68,19 @@ def __init__(
         # Construction of hidden layers
         # H_i: ℝ^{d_i} → ℝ^{d_{i+1}}
         # Sequential transformation φ(x) = Dropout(ReLU(LayerNorm(Wx + b)))
+        logger.debug("Constructing network layers")
         self.layers = nn.ModuleList()
         prev_dim = in_features
 
-        for dim in hidden_dims:
-
-            # Affine transformation followed by distribution normalisation
+        for i, dim in enumerate(hidden_dims):
+            logger.debug(f"Building layer {i+1}: {prev_dim} → {dim}")
             layer_block = []
             layer_block.append(nn.Linear(prev_dim, dim))
 
             if use_layer_norm:
+                logger.debug(f"Adding LayerNorm for dimension {dim}")
                 layer_block.append(nn.LayerNorm(dim))
 
-            # Non-linear activation and stochastic regularisation
             layer_block.extend([
                 nn.ReLU(),
                 nn.Dropout(dropout)
@@ -83,14 +94,17 @@ def __init__(
         # Projection mapping P: ℝ^d → ℝ^{m×t} for temporal quantile predictions
         if quantile_output and num_forecast_steps:
             final_out_features = out_features * num_forecast_steps
+            logger.debug(f"Configuring for quantile output with {num_forecast_steps} steps")
         else:
             final_out_features = out_features
-            
+
+        logger.debug(f"Creating output layer: {prev_dim} → {final_out_features}")
         self.output_layer = nn.Linear(prev_dim, final_out_features)
 
         # Output activation definition
         # ψ: ℝ^m → [0,1]^m
         if output_activation == "softmax":
+            logger.debug(f"Setting output activation: {output_activation}")
             self.output_activation = nn.Softmax(dim=-1)
         elif output_activation == "sigmoid":
             self.output_activation = nn.Sigmoid()
@@ -100,6 +114,7 @@ def __init__(
         # Optional layer norm and residual projection
         # g: ℝ^n → ℝ^m
         if use_residual:
+            logger.debug("Initialising residual connection components")
             if quantile_output and num_forecast_steps:
                 self.residual_norm = nn.LayerNorm(out_features)
             else:
@@ -108,10 +123,11 @@ def __init__(
             self.residual_proj = nn.Linear(in_features, out_features)
 
     def reshape_quantile_output(self, x: torch.Tensor) -> torch.Tensor:
-
-        # Reshape output for quantile predictions
+        logger.debug(f"Input shape before reshape: {x.shape}")
         if self.quantile_output and self.num_forecast_steps:
-            return x.reshape(x.shape[0], self.num_forecast_steps, -1)
+            reshaped = x.reshape(x.shape[0], self.num_forecast_steps, -1)
+            logger.debug(f"Reshaped output shape: {reshaped.shape}")
+            return reshaped
         return x
 
     def forward(
@@ -120,47 +136,59 @@ def forward(
         return_intermediates: bool = False
     ) -> Union[torch.Tensor, tuple]:
 
+        logger.info("Starting DynamicOutputNetwork forward pass")
+
         # Forward pass for dynamic output network
         # Handle dict input
         # Concatenate multimodal inputs if dict provided
         if isinstance(x, dict):
+            logger.debug(f"Processing dictionary input with keys: {list(x.keys())}")
             x = torch.cat(list(x.values()), dim=-1)
+            logger.debug(f"Concatenated input shape: {x.shape}")
 
         intermediates = []
         residual = x
 
         # Process through hidden layers
-        for layer in self.layers:
+        for i, layer in enumerate(self.layers):
+            logger.debug(f"Processing layer {i+1}, input shape: {x.shape}")
             x = layer(x)
             if return_intermediates:
                 intermediates.append(x)
 
         # Output transform, reshape and apply residual connection
+        logger.debug(f"Applying output layer to shape: {x.shape}")
         x = self.output_layer(x)        
         x = self.reshape_quantile_output(x)    
 
         if self.use_residual:
+            logger.debug("Applying residual connection")
+
             # Apply residual projection transformation
             projected_residual = self.residual_proj(residual)
             if self.quantile_output and self.num_forecast_steps:
+                logger.debug("Processing quantile output with residual")
 
                 # Apply residual mapping followed by normalisation
                 projected_residual = projected_residual.reshape(x.shape[0], x.shape[2])
 
                 # Collapse temporal dimensions for normalisation 
                 # ℝ^{B×T×F} → ℝ^{BT×F}
                 x = x.reshape(-1, x.shape[2])
+                logger.debug(f"Reshaped for residual: {x.shape}")
                 x = self.residual_norm(x + projected_residual.repeat(self.num_forecast_steps, 1))
 
                 # Restore tensor dimensionality
                 # ℝ^{BT×F} → ℝ^{B×T×F}
                 x = x.reshape(-1, self.num_forecast_steps, self.out_features)
+                logger.debug(f"Final shape after residual: {x.shape}")
             else:
                 x = self.residual_norm(x + projected_residual)
 
         # Apply output activation
         # Non-linear transformation ψ
         if self.output_activation:
+            logger.debug(f"Applying output activation: {type(self.output_activation).__name__}")
             x = self.output_activation(x)
 
         if return_intermediates:
@@ -180,7 +208,10 @@ def __init__(
         hidden_dims: Optional[List[int]] = None,
         dropout: float = 0.1
     ):
-        
+
+        logger.info(f"Initialising QuantileOutputNetwork with in_features={in_features}, num_quantiles={num_quantiles}")
+        logger.debug(f"Forecast steps: {num_forecast_steps}, hidden_dims: {hidden_dims}")
+
         # Initialisation of quantile output network
         super().__init__(
             in_features=in_features,