# PyTorch Training Framework Refactoring Guide

This guide documents all the improvements made to the PyTorch training framework, explaining the motivation behind each change and how to use the new features.

## Overview of Changes

### 1. **Modular Training Architecture**

**Location**: `training/` module

**What changed**:
- Created a reusable `Trainer` class that encapsulates training logic
- Separated concerns: model training, metric calculation, logging, and checkpointing
- Implemented a hook system for extensibility

**Why**:
- **Modularity**: Training logic is now independent of model specifics
- **Reusability**: Same trainer can be used for different models
- **Extensibility**: Easy to add new features without modifying core logic
- **Maintainability**: Cleaner code structure with single responsibility principle

**How to use**:
```python
from training import Trainer, TrainingHooks
from training.hooks import WandbHook, CheckpointHook

# Create hooks
hooks = TrainingHooks()
hooks.add_hook(WandbHook(config=args))
hooks.add_hook(CheckpointHook(save_dir=output_dir))

# Create trainer
trainer = Trainer(
    model=model,
    optimizer=optimizer,
    config=args,
    train_loader=train_loader,
    val_loader=val_loader,
    hooks=hooks
)

# Start training
trainer.train()
```

### 2. **Configurable Learning Rate Schedulers**

**Location**: `training/schedulers.py`

**What changed**:
- Created `SchedulerFactory` for centralized scheduler creation
- Support for 7 different scheduler types
- Fixed the CosineAnnealingLR float mismatch error
- Added warmup support for any scheduler

**Why**:
- **Flexibility**: Easy to experiment with different schedulers
- **Type Safety**: Automatic type conversion prevents errors
- **Configuration**: Schedulers configured via YAML, no code changes needed
- **Best Practices**: Implements modern training techniques like warmup

**Supported schedulers**:
1. **MultiStepLR**: Decay at specific milestones
2. **CosineAnnealingLR**: Smooth cosine annealing (fixed float issue)
3. **StepLR**: Simple step decay
4. **ReduceLROnPlateau**: Adaptive based on validation metric
5. **OneCycleLR**: Super-convergence training
6. **ExponentialLR**: Exponential decay
7. **CyclicLR**: Cyclic learning rates

**Configuration example**:
```yaml
# In your config YAML file
scheduler:
  type: cosine
  params:
    T_max: 100  # Will be converted to int automatically
    eta_min: 1e-6

# Optional warmup
warmup_epochs: 5
```

### 3. **Optimized Training Engine**

**Location**: `engine/crog_engine_optimized.py`

**What changed**:
- Removed code duplication between train/validate functions
- Created `MetricCalculator` class for centralized metric computation
- Batch processing for validation transforms
- More efficient memory management

**Why**:
- **Performance**: Reduced redundant calculations
- **Memory**: Better GPU memory utilization
- **Maintainability**: Single source of truth for metrics
- **Clarity**: Cleaner code structure

**Key optimizations**:
```python
# Before: Repeated code for each metric
iou = calculate_iou(pred, target)
precision = calculate_precision(pred, target)

# After: Centralized calculation
metric_calc = MetricCalculator()
metrics = metric_calc.calculate_batch_metrics(pred, target)
```

### 4. **Enhanced WandB Integration**

**Location**: `utils/wandb_utils.py`, `training/hooks.py`

**What changed**:
- Removed hardcoded API keys (security fix)
- Created `WandBLogger` with proper metric definitions
- Fixed metric visualization issues
- Added comprehensive metric tracking

**Why**:
- **Security**: API keys should never be hardcoded
- **Visualization**: Proper metric definitions ensure correct graphs
- **Completeness**: Track all relevant metrics including gradients
- **Debugging**: Better insight into training process

**New features**:
1. **Automatic metric namespacing**: training/, validation/, lr/
2. **Custom metric definitions**: Proper step tracking
3. **Gradient tracking**: Monitor gradient norms
4. **Learning rate per parameter group**: Track different LRs

**Setup**:
```bash
# Set API key securely
export WANDB_API_KEY=your_api_key_here

# Or use .netrc file (recommended)
echo "machine api.wandb.ai
  login user
  password your_api_key_here" >> ~/.netrc
chmod 600 ~/.netrc
```

### 5. **Hook System for Extensibility**

**Location**: `training/hooks.py`

**What changed**:
- Implemented hook system with base `Hook` class
- Created specialized hooks: WandBHook, CheckpointHook, MetricLoggerHook
- Hooks called at specific points in training

**Why**:
- **Separation of Concerns**: Logging, checkpointing separate from training
- **Extensibility**: Easy to add new functionality
- **Reusability**: Hooks can be shared across projects
- **Testing**: Easier to test individual components

**Creating custom hooks**:
```python
class CustomHook(Hook):
    def on_epoch_end(self, epoch: int, **kwargs):
        # Your custom logic here
        print(f"Epoch {epoch} completed!")
```

## Migration Guide

### For Existing Scripts

1. **Update scheduler configuration**:
   ```yaml
   # Old
   milestones: [30, 60]
   lr_decay: 0.1
   
   # New
   scheduler:
     type: multistep
     params:
       milestones: [30, 60]
       gamma: 0.1
   ```

2. **Update WandB initialization**:
   ```python
   # Old
   os.environ["WANDB_API_KEY"] = 'hardcoded_key'
   wandb.init(...)
   
   # New
   # Set WANDB_API_KEY environment variable
   wandb_logger = WandBLogger(...)
   ```

3. **Use the refactored trainer** (optional):
   ```python
   # See train_network_refactored.py for complete example
   ```

### For New Projects

1. Use `train_network_refactored.py` as a template
2. Configure schedulers in YAML using `config/scheduler_examples.yaml`
3. Set up WandB following `docs/wandb_setup.md`
4. Extend functionality using the hook system

## Benefits Summary

1. **Modularity**: Clean separation of concerns
2. **Flexibility**: Easy configuration without code changes
3. **Performance**: Optimized training loops
4. **Debugging**: Better logging and visualization
5. **Maintainability**: Cleaner, more organized code
6. **Extensibility**: Hook system for custom features
7. **Best Practices**: Implements modern deep learning techniques

## Backward Compatibility

All changes maintain backward compatibility:
- Old config files still work (auto-converted to new format)
- Existing models and datasets unchanged
- Training results remain consistent

## Future Improvements

Potential enhancements:
1. Add more schedulers (polynomial decay, cosine restarts)
2. Implement automatic hyperparameter tuning
3. Add distributed training optimizations
4. Create visualization dashboard
5. Add experiment tracking beyond WandB