minor fixed on src/main/parser.py

src/main.py

@@ -32,13 +32,14 @@ Dataset Requirements:
             └── ...
     
     Each .mat file must contain variable 'H' with shape [subcarriers, symbols, 3]:
-    - H[:, :, 0]: Ground truth channel
-    - H[:, :, 1]: LS channel estimate with zeros for non-pilot positions
-    - H[:, :, 2]: Unused (reserved)
+    - H[:, :, 0]: Ground truth channel (complex-valued channel matrix)
+    - H[:, :, 1]: LS channel estimate with zeros for non-pilot positions (complex-valued) - used as input to models
+    - H[:, :, 2]: Bilinear interpolated LS channel estimate (complex-valued) - available but currently unused
 """
 
 import logging
 import sys
+from datetime import datetime
 from pathlib import Path
 
 from src.main.parser import parse_arguments
@@ -47,18 +48,26 @@ from src.config import load_config
 from src.config.schemas import ModelConfig
 
 
-def setup_logging(log_level: str) -> None:
+def setup_logging(log_level: str, log_dir: Path, exp_id: str) -> None:
     """Set up logging configuration.
     
     Args:
         log_level: Logging level string (DEBUG, INFO, WARNING, ERROR, CRITICAL)
+        log_dir: Directory path for log files
+        exp_id: Experiment identifier for log file naming
     """
+    # Create logs directory if it doesn't exist
+    log_dir.mkdir(parents=True, exist_ok=True)
+    
+    # Create log file path using exp_id for easy matching
+    log_file = log_dir / f"training_{exp_id}.log"
+    
     logging.basicConfig(
         level=getattr(logging, log_level.upper()),
         format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
         handlers=[
             logging.StreamHandler(sys.stdout),
-            logging.FileHandler('training.log')
+            logging.FileHandler(log_file)
         ]
     )
 
@@ -70,7 +79,7 @@ def main() -> None:
         args = parse_arguments()
         
         # Set up logging
-        setup_logging(args.python_log_level)
+        setup_logging(args.python_log_level, args.python_log_dir, args.exp_id)
         logger = logging.getLogger(__name__)
         
         logger.info("Starting OFDM channel estimation model training")
@@ -86,13 +95,35 @@ def main() -> None:
             args.model_config_path
         )
         
+        # Validate model type consistency
+        expected_model_types = {
+            "linear": "linear",
+            "fortitran": "fortitran", 
+            "adafortitran": "adafortitran"
+        }
+        
+        if args.model_name not in expected_model_types:
+            raise ValueError(f"Unknown model name: {args.model_name}. Expected one of: {list(expected_model_types.keys())}")
+        
+        if model_config.model_type != expected_model_types[args.model_name]:
+            raise ValueError(f"Model type mismatch: config specifies '{model_config.model_type}' but model name is '{args.model_name}'")
+        
         logger.info("Configuration loaded successfully")
         logger.info(f"OFDM dimensions: {system_config.ofdm.num_scs} subcarriers x {system_config.ofdm.num_symbols} symbols")
         logger.info(f"Pilot dimensions: {system_config.pilot.num_scs} subcarriers x {system_config.pilot.num_symbols} symbols")
+        
+        # Log model-specific information
         if model_config.model_type == "linear":
             logger.info(f"Linear model with device: {model_config.device}")
+        elif model_config.model_type == "fortitran":
+            logger.info(f"FortiTran model: {model_config.num_layers} layers, {model_config.model_dim} dimensions")
+            logger.info(f"Channel adaptation: disabled")
+        elif model_config.model_type == "adafortitran":
+            logger.info(f"AdaFortiTran model: {model_config.num_layers} layers, {model_config.model_dim} dimensions")
+            logger.info(f"Channel adaptation: enabled")
+            logger.info(f"Adaptive token length: {model_config.adaptive_token_length}")
         else:
-            logger.info(f"Model architecture: {model_config.num_layers} layers, {model_config.model_dim} dimensions")
+            logger.warning(f"Unknown model type: {model_config.model_type}")
         
         # Start training
         logger.info("Initializing training...")

src/main/parser.py

@@ -9,18 +9,10 @@ of training runs.
 
 from pathlib import Path
 import argparse
-from enum import Enum
 from pydantic import BaseModel, Field, model_validator
 from typing import Self
 
 
-class LossType(Enum):
-    """Enumeration of supported loss functions."""
-    MSE = "mse"
-    MAE = "mae"
-    HUBER = "huber"
-
-
 class TrainingArguments(BaseModel):
     """Container for OFDM model training arguments.
 
@@ -29,7 +21,7 @@ class TrainingArguments(BaseModel):
 
     Attributes:
         # Model Configuration
-        model_name: Supports Linear, AdaFortiTran, or FortiTran training
+        model_name: Supports linear, adafortitran, or fortitran training
         system_config_path: Path to OFDM system configuration file
         model_config_path: Path to model configuration file
 
@@ -42,17 +34,15 @@ class TrainingArguments(BaseModel):
         exp_id: Experiment identifier string
         python_log_level: Logging verbosity level
         tensorboard_log_dir: Directory for tensorboard logs
+        python_log_dir: Directory for python logging files
 
         # Training Hyperparameters
         batch_size: Number of samples per batch
         lr: Learning rate for optimizer
         max_epoch: Maximum number of training epochs
         patience: Early stopping patience in epochs
-        loss_type: Type of loss function to use
-        return_type: Type of data to return from dataset
 
-        # Hardware & Evaluation
-        cuda: CUDA device index
+        # Evaluation
         test_every_n: Number of epochs between test evaluations
     """
 
@@ -70,17 +60,15 @@ class TrainingArguments(BaseModel):
     exp_id: str = Field(..., description="Experiment identifier for log folder naming")
     python_log_level: str = Field(default="INFO", description="Logger level for python logging module")
     tensorboard_log_dir: Path = Field(default=Path("runs"), description="Directory for tensorboard logs")
+    python_log_dir: Path = Field(default=Path("logs"), description="Directory for python logging files")
 
     # Training Hyperparameters
     batch_size: int = Field(default=64, gt=0, description="Training batch size")
     lr: float = Field(default=1e-3, gt=0, description="Initial learning rate")
     max_epoch: int = Field(default=10, gt=0, description="Maximum number of training epochs")
     patience: int = Field(default=3, gt=0, description="Early stopping patience (epochs)")
-    loss_type: LossType = Field(default=LossType.MSE, description="Loss function type")
-    return_type: str = Field(default="complex", description="Type of data to return from dataset")
 
-    # Hardware & Evaluation
-    cuda: int = Field(default=0, ge=0, description="CUDA device index (0 for single GPU)")
+    # Evaluation
     test_every_n: int = Field(default=10, gt=0, description="Test model every N epochs")
 
     @model_validator(mode='after')
@@ -133,8 +121,8 @@ def parse_arguments() -> TrainingArguments:
         '--model_name',
         type=str,
         required=True,
-        choices=['Linear', 'AdaFortiTran', 'FortiTran'],
-        help='Model type to train (Linear, AdaFortiTran, or FortiTran)'
+        choices=['linear', 'adafortitran', 'fortitran'],
+        help='Model type to train (linear, adafortitran, or fortitran)'
     )
     required.add_argument(
         '--system_config_path',
@@ -187,6 +175,12 @@ def parse_arguments() -> TrainingArguments:
         default="runs",
         help='Directory for tensorboard logs'
     )
+    optional.add_argument(
+        '--python_log_dir',
+        type=Path,
+        default="logs",
+        help='Directory for python logging files'
+    )
     optional.add_argument(
         '--test_every_n',
         type=int,
@@ -211,55 +205,16 @@ def parse_arguments() -> TrainingArguments:
         default=64,
         help='Training batch size'
     )
-    optional.add_argument(
-        '--cuda',
-        type=int,
-        default=0,
-        help='CUDA device index (0 for single GPU)'
-    )
+
     optional.add_argument(
         '--lr',
         type=float,
         default=1e-3,
         help='Initial learning rate'
     )
-    optional.add_argument(
-        '--loss_type',
-        type=str,
-        default="mse",
-        choices=['mse', 'mae', 'huber'],
-        help='Loss function type'
-    )
-    optional.add_argument(
-        '--return_type',
-        type=str,
-        default="complex",
-        choices=['complex', 'real'],
-        help='Type of data to return from dataset'
-    )
 
-    args = parser.parse_args()
 
-    # Convert loss_type string to enum
-    loss_type = LossType(args.loss_type)
+    args = parser.parse_args()
 
     # Create and validate TrainingArguments
-    return TrainingArguments(
-        model_name=args.model_name,
-        system_config_path=args.system_config_path,
-        model_config_path=args.model_config_path,
-        train_set=args.train_set,
-        val_set=args.val_set,
-        test_set=args.test_set,
-        exp_id=args.exp_id,
-        python_log_level=args.python_log_level,
-        tensorboard_log_dir=args.tensorboard_log_dir,
-        batch_size=args.batch_size,
-        lr=args.lr,
-        max_epoch=args.max_epoch,
-        patience=args.patience,
-        loss_type=loss_type,
-        return_type=args.return_type,
-        cuda=args.cuda,
-        test_every_n=args.test_every_n
-    )
+    return TrainingArguments(**vars(args))

src/main/train_helpers.py

@@ -120,7 +120,7 @@ def eval_model(
             output = _compute_loss(estimated_channel, ideal_channel, loss_fn)
             val_loss += (2 * output.item() * batch[0].size(0))
 
-    val_loss /= len(eval_dataloader.dataset)
+    val_loss /= sum(len(batch[0]) for batch in eval_dataloader)
     return val_loss
 
 
@@ -206,7 +206,7 @@ def train_epoch(
         train_loss += (2 * output.item() * batch[0].size(0))
 
     scheduler.step()
-    train_loss /= len(train_dataloader.dataset)
+    train_loss /= sum(len(batch[0]) for batch in train_dataloader)
     return train_loss
 
 
@@ -225,20 +225,17 @@ def _forward_pass(batch: BatchType, model: nn.Module) -> Tuple[ComplexTensor, Co
         Tuple of (processed_estimated_channel, ideal_channel)
 
     Raises:
-        ValueError: If model name is not recognized
+        ValueError: If model type is not recognized
     """
     estimated_channel, ideal_channel, meta_data = batch
 
-    if model.name in ["fortitran", "MMSE"]:
-        h_est_re = model(torch.real(estimated_channel))
-        h_est_im = model(torch.imag(estimated_channel))
-        estimated_channel = torch.complex(h_est_re, h_est_im)
-    elif model.name == "adafortitran":
-        h_est_re = model(torch.real(estimated_channel), meta_data)
-        h_est_im = model(torch.imag(estimated_channel), meta_data)
-        estimated_channel = torch.complex(h_est_re, h_est_im)
+    # All models now handle complex input directly
+    if hasattr(model, 'use_channel_adaptation') and model.use_channel_adaptation:
+        # AdaFortiTran uses meta_data for channel adaptation
+        estimated_channel = model(estimated_channel, meta_data)
     else:
-        raise ValueError(f"Unknown model type: {model.name}")
+        # Linear and FortiTran models don't use meta_data
+        estimated_channel = model(estimated_channel)
 
     return estimated_channel, ideal_channel.to(model.device)
 

src/main/trainer.py

@@ -81,7 +81,7 @@ class ModelTrainer:
         self.system_config = system_config
         self.model_config = model_config
         self.args = args
-        self.device = torch.device(f"cuda:{args.cuda}")
+        self.device = torch.device(model_config.device)
         self.writer = self._setup_tensorboard()
         self.logger = logging.getLogger(__name__)
 
@@ -120,14 +120,12 @@ class ModelTrainer:
         Returns:
             Initialized model instance of the specified type
         """
-        if self.args.model_name == "linear":
-            model = LinearEstimator(self.system_config, self.model_config)
-        elif self.args.model_name == "adafortitran":
-            model = AdaFortiTranEstimator(self.system_config, self.model_config)
-        elif self.args.model_name == "fortitran":
-            model = FortiTranEstimator(self.system_config, self.model_config)
-        else:
-            raise ValueError(f"Unknown model name: {self.args.model_name}")
+        if self.args.model_name not in self.MODEL_REGISTRY:
+            raise ValueError(f"Unknown model name: {self.args.model_name}. Available: {list(self.MODEL_REGISTRY.keys())}")
+        
+        model_class = self.MODEL_REGISTRY[self.args.model_name]
+        model = model_class(self.system_config, self.model_config)
+        
         num_params, model_summary = get_model_details(model)
         self.logger.info("\n" + model_summary)
         self.logger.info(f"Model name: {self.args.model_name} | Number of parameters: {num_params}")
@@ -280,20 +278,15 @@ class ModelTrainer:
 
     def _forward_pass(self, batch, model):
         estimated_channel, ideal_channel, meta_data = batch
-        if isinstance(model, FortiTranEstimator):
-            h_est_re = model(torch.real(estimated_channel))
-            h_est_im = model(torch.imag(estimated_channel))
-            estimated_channel = torch.complex(h_est_re, h_est_im)
-        elif isinstance(model, AdaFortiTranEstimator):
-            h_est_re = model(torch.real(estimated_channel), meta_data)
-            h_est_im = model(torch.imag(estimated_channel), meta_data)
-            estimated_channel = torch.complex(h_est_re, h_est_im)
-        elif isinstance(model, LinearEstimator):
-            h_est_re = model(torch.real(estimated_channel))
-            h_est_im = model(torch.imag(estimated_channel))
-            estimated_channel = torch.complex(h_est_re, h_est_im)
+        
+        # All models now handle complex input directly
+        if isinstance(model, AdaFortiTranEstimator):
+            # AdaFortiTran uses meta_data for channel adaptation
+            estimated_channel = model(estimated_channel, meta_data)
         else:
-            raise ValueError(f"Unknown model type: {type(model)}")
+            # Linear and FortiTran models don't use meta_data
+            estimated_channel = model(estimated_channel)
+            
         return estimated_channel, ideal_channel.to(model.device)
 
     def _train_epoch(self):