REFACTORING_GUIDE.md

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:

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:

# 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:

# 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:

# 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:

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:

   # Old
   milestones: [30, 60]
   lr_decay: 0.1
   
   # New
   scheduler:
     type: multistep
     params:
       milestones: [30, 60]
       gamma: 0.1
   

2. Update WandB initialization:

   # 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):

   # 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