Logging API

This section documents the centralized logging system.

Base Logger

Base functionality for all loggers:

from jepa.loggers.base_logger import BaseLogger

class CustomLogger(BaseLogger):
    def __init__(self, config):
        super().__init__(config)
        # Custom initialization
    
    def log_metrics(self, metrics, step=None):
        # Custom logging implementation
        pass

Logger Registry

Manage available logger types:

from jepa.loggers.base_logger import LoggerRegistry

# Register custom logger
LoggerRegistry.register("custom", CustomLogger)

# Create logger by name
logger = LoggerRegistry.create("wandb", config)

Console Logger

Simple console-based logging:

from jepa.loggers import ConsoleLogger

logger = ConsoleLogger({
    'level': 'INFO',
    'format': '%(asctime)s - %(levelname)s - %(message)s',
    'colored': True
})

logger.log_metrics({'loss': 0.5, 'accuracy': 0.85}, step=100)

Weights & Biases Logger

Integration with Weights & Biases:

from jepa.loggers import WandbLogger

logger = WandbLogger({
    'project': 'jepa-experiments',
    'entity': 'my-team',
    'name': 'experiment-1',
    'tags': ['baseline', 'transformer'],
    'notes': 'Initial experiment with transformer encoder'
})

# Log metrics
logger.log_metrics({'train/loss': 0.5, 'train/lr': 1e-4}, step=100)

# Log model
logger.log_model(model, 'model_checkpoint')

TensorBoard Logger

TensorBoard integration for visualization:

from jepa.loggers import TensorBoardLogger

logger = TensorBoardLogger({
    'log_dir': 'logs/tensorboard',
    'log_frequency': 10,
    'log_images': True,
    'log_histograms': True
})

# Log scalars
logger.log_metrics({'loss': 0.5}, step=100)

# Log images
logger.log_image('input_samples', images, step=100)

# Log histograms
logger.log_histogram('model.encoder.weight', model.encoder.weight, step=100)

Multi Logger

Combine multiple logging backends:

from jepa.loggers import MultiLogger

# List-based configuration
logger = MultiLogger([
    ('console', {'level': 'INFO'}),
    ('wandb', {'project': 'jepa'}),
    ('tensorboard', {'log_dir': 'logs'})
])

# Dictionary-based configuration
logger = MultiLogger({
    'backends': ['console', 'wandb', 'tensorboard'],
    'console': {'level': 'INFO'},
    'wandb': {'project': 'jepa'},
    'tensorboard': {'log_dir': 'logs'}
})

# Logs to all backends simultaneously
logger.log_metrics({'loss': 0.5}, step=100)

Logger Factory

Create loggers from configuration:

from jepa.loggers import create_logger

# From dictionary config
config = {
    'backends': ['wandb', 'tensorboard'],
    'wandb': {'project': 'jepa-experiments'},
    'tensorboard': {'log_dir': 'logs/tb'}
}

logger = create_logger(config)

# From YAML config
logger = create_logger("config/logging_config.yaml")

Logging Utilities

Metric Formatting

Utilities for metric processing:

from jepa.loggers.utils import format_metrics, add_prefix

# Format metrics for display
formatted = format_metrics({'loss': 0.123456}, precision=4)
# Output: {'loss': '0.1235'}

# Add prefix to metrics
prefixed = add_prefix({'loss': 0.5, 'acc': 0.9}, 'train')
# Output: {'train/loss': 0.5, 'train/acc': 0.9}

Logging Decorators

Decorators for automatic logging:

from jepa.loggers.decorators import log_execution_time, log_exceptions

@log_execution_time(logger)
@log_exceptions(logger)
def training_step(batch):
    # Function execution time and exceptions are automatically logged
    return process_batch(batch)

Advanced Logging Features

Structured Logging

Log structured data with schema validation:

from jepa.loggers.structured import StructuredLogger

logger = StructuredLogger({
    'schema': {
        'metrics': {'type': 'dict', 'required': True},
        'metadata': {'type': 'dict', 'required': False}
    }
})

logger.log_structured({
    'metrics': {'loss': 0.5, 'accuracy': 0.9},
    'metadata': {'epoch': 10, 'batch_size': 64}
})

Async Logging

Non-blocking logging for high-throughput scenarios:

from jepa.loggers.async_logger import AsyncLogger

logger = AsyncLogger({
    'backend': 'wandb',
    'buffer_size': 1000,
    'flush_interval': 30  # seconds
})

# Non-blocking log calls
logger.log_metrics({'loss': 0.5}, step=100)

Conditional Logging

Log based on conditions:

from jepa.loggers.conditional import ConditionalLogger

logger = ConditionalLogger({
    'base_logger': wandb_logger,
    'conditions': {
        'min_step': 100,  # Only log after step 100
        'max_frequency': 0.1,  # Max 10% of calls
        'level_filter': 'INFO'  # Only INFO and above
    }
})

Configuration

Logging Configuration

Configure logging through YAML:

logging:
  level: "INFO"
  backends: ["wandb", "tensorboard", "console"]
  
  # Backend-specific settings
  wandb:
    project: "jepa-experiments"
    entity: "my-team"
    tags: ["baseline", "transformer"]
    log_frequency: 10
    log_gradients: false
    log_model: true
  
  tensorboard:
    log_dir: "logs/tensorboard"
    log_frequency: 10
    log_images: true
    log_histograms: true
  
  console:
    level: "INFO"
    format: "%(asctime)s - %(levelname)s - %(message)s"
    colored: true
  
  # Global settings
  log_frequency: 10
  save_frequency: 50

Custom Logger Development

Creating Custom Loggers

from jepa.loggers.base_logger import BaseLogger

class CustomAPILogger(BaseLogger):
    def __init__(self, config):
        super().__init__(config)
        self.api_endpoint = config.get('api_endpoint')
        self.api_key = config.get('api_key')
        self.session = requests.Session()
    
    def log_metrics(self, metrics, step=None):
        payload = {
            'metrics': metrics,
            'step': step,
            'timestamp': time.time()
        }
        
        response = self.session.post(
            self.api_endpoint,
            json=payload,
            headers={'Authorization': f'Bearer {self.api_key}'}
        )
        
        if not response.ok:
            self.logger.warning(f"Failed to log metrics: {response.status_code}")
    
    def finish(self):
        self.session.close()

# Register the custom logger
from jepa.loggers.base_logger import LoggerRegistry
LoggerRegistry.register('custom_api', CustomAPILogger)

Logger Middleware

from jepa.loggers.middleware import LoggerMiddleware

class MetricFilterMiddleware(LoggerMiddleware):
    def __init__(self, filter_patterns):
        self.filter_patterns = filter_patterns
    
    def process_metrics(self, metrics, step=None):
        filtered_metrics = {}
        for key, value in metrics.items():
            if any(pattern in key for pattern in self.filter_patterns):
                filtered_metrics[key] = value
        return filtered_metrics, step

# Apply middleware
logger = WandbLogger(config)
logger.add_middleware(MetricFilterMiddleware(['train', 'val']))

Integration Examples

Basic Logging Setup

from jepa.loggers import create_logger

# Simple console logging
logger = create_logger({'backends': ['console']})

# Multi-backend logging
logger = create_logger({
    'backends': ['wandb', 'tensorboard', 'console'],
    'wandb': {'project': 'my-project'},
    'tensorboard': {'log_dir': 'logs'},
    'console': {'level': 'INFO'}
})

# Use in training
for epoch in range(epochs):
    for batch_idx, batch in enumerate(dataloader):
        loss = training_step(batch)
        
        if batch_idx % log_frequency == 0:
            logger.log_metrics({
                'train/loss': loss,
                'train/epoch': epoch,
                'train/step': global_step
            }, step=global_step)

Advanced Logging

To integrate custom logic, subclass your logger or wrap it in your own helper and pass it to the trainer:

from jepa.loggers import BaseLogger
from jepa.trainer import create_trainer

class AggregatingLogger(BaseLogger):
    def __init__(self):
        super().__init__({"enabled": True})
        self.history = []

    def log_metrics(self, metrics, step=None):
        enriched = {"step": step, **metrics}
        self.history.append(enriched)
        print("LOG", enriched)

    def log_hyperparameters(self, hyperparams):
        print("HYPERPARAMS", hyperparams)

    def save_artifact(self, file_path, artifact_name=None):
        pass

    def finish(self):
        print("Finished logging", len(self.history), "records")

logger = AggregatingLogger()
trainer = create_trainer(model, logger=logger)
trainer.train(train_loader, num_epochs=5)  # `train_loader` defined elsewhere

# Or use the shortcut to enable a single backend
wandb_trainer = create_trainer(
    model,
    logger="wandb",
    logger_project="jepa-demo",
    logger_run_name="experiment-001",
)

For more examples and detailed usage, see the Logging Examples and Training Guide.