training/__init__.py

"""
Training Module
==============
This module provides a modular training framework for PyTorch models.

Key components:
1. **Trainer**: Base training class with hooks support
2. **SchedulerFactory**: Factory for creating LR schedulers
3. **TrainingHooks**: Hook system for extensible training
4. **GraspValidator**: Specialized validator for grasp predictions
5. **Enhanced Hooks**: Extended hooks with additional metrics
"""

# Modular training framework components
from .trainer import Trainer
from .schedulers import SchedulerFactory
from .hooks import TrainingHooks, Hook, WandbHook, CheckpointHook, MetricLoggerHook
from .grasp_validator import GraspValidator

# Enhanced hooks with additional metrics
try:
    from .enhanced_hooks import EnhancedWandbHook, EnhancedCheckpointHook, ValidationTimerHook
    __all__ = [
        'Trainer',
        'SchedulerFactory', 
        'TrainingHooks',
        'Hook',
        'WandbHook',
        'CheckpointHook',
        'MetricLoggerHook',
        'GraspValidator',
        'EnhancedWandbHook',
        'EnhancedCheckpointHook',
        'ValidationTimerHook'
    ]
except ImportError:
    # Fall back to basic hooks if enhanced hooks not available
    __all__ = [
        'Trainer',
        'SchedulerFactory', 
        'TrainingHooks',
        'Hook',
        'WandbHook',
        'CheckpointHook',
        'MetricLoggerHook',
        'GraspValidator'
    ]