Spaces:
Running on Zero
Running on Zero
| import wandb | |
| import torch | |
| import os | |
| import time | |
| from typing import Dict, Any, Optional | |
| def init_wandb(project_name: str, config: Dict[str, Any], run_name: Optional[str] = None) -> wandb.run: | |
| # Generate run name if not provided | |
| if run_name is None: | |
| run_name = f"{project_name}_{time.strftime('%Y%m%d_%H%M%S')}" | |
| # Initialize wandb | |
| run = wandb.init( | |
| project=project_name, | |
| config=config, | |
| name=run_name, | |
| tags=[project_name, "training"] | |
| ) | |
| print(f"🚀 W&B run initialized: {run.name}") | |
| print(f"📊 Project: {project_name}") | |
| print(f"🔗 View at: {run.url}") | |
| return run | |
| def log_training_metrics(step: int, metrics: Dict[str, float], prefix: str = "train"): | |
| # Add prefix to metric names | |
| prefixed_metrics = {f"{prefix}/{k}": (v.item() if hasattr(v, "item") else float(v)) for k, v in metrics.items()} | |
| wandb.log(prefixed_metrics, step=step) | |
| def log_learning_rate(optimizer: torch.optim.Optimizer, step: int): | |
| for i, param_group in enumerate(optimizer.param_groups): | |
| wandb.log({ | |
| f"learning_rate/group_{i}": float(param_group['lr']), | |
| }, step=step) | |
| def log_codebook_usage(codebook_usage: float, step: int, model_name: str = "model"): | |
| wandb.log({ | |
| f"{model_name}/codebook_usage": float(codebook_usage), | |
| }, step=step) | |
| def log_action_distribution(action_indices: torch.Tensor, step: int, n_actions: int): | |
| # Convert to CPU and numpy | |
| flat = action_indices.detach().cpu().reshape(-1) | |
| # Compute distribution | |
| counts = torch.bincount(flat.long(), minlength=n_actions).float() | |
| probs = counts / counts.sum().clamp_min(1) | |
| wandb.log({ | |
| "action_distribution": wandb.Histogram(flat.numpy()), | |
| "action_entropy": float(-(probs * (probs + 1e-8).log()).sum()), | |
| "unique_actions": int((counts > 0).sum().item()), | |
| }, step=step) | |
| def log_system_metrics(step: int): | |
| if torch.cuda.is_available(): | |
| wandb.log({ | |
| "system/gpu_memory_allocated": float(torch.cuda.memory_allocated() / 1024**3), # GB | |
| "system/gpu_memory_reserved": float(torch.cuda.memory_reserved() / 1024**3), # GB | |
| }, step=step) | |
| def finish_wandb(): | |
| """Finish the W&B run""" | |
| if wandb.run is not None: | |
| wandb.finish() | |
| def create_wandb_config(args, model_config: Dict[str, Any]) -> Dict[str, Any]: | |
| config = { | |
| # Training parameters | |
| "batch_size": args.batch_size, | |
| "n_updates": args.n_updates, | |
| "learning_rate": args.learning_rate, | |
| "log_interval": getattr(args, 'log_interval', 100), | |
| # Dataset parameters | |
| "dataset": getattr(args, 'dataset', 'SONIC'), | |
| "context_length": getattr(args, 'context_length', 4), | |
| # Model architecture | |
| "model_architecture": model_config, | |
| # System parameters | |
| "device": str(torch.device("cuda" if torch.cuda.is_available() else "cpu")), | |
| "timestamp": time.strftime("%Y%m%d_%H%M%S"), | |
| } | |
| # Add any additional args that might exist | |
| for attr in dir(args): | |
| if not attr.startswith('_') and not callable(getattr(args, attr)): | |
| if attr not in config: | |
| config[attr] = getattr(args, attr) | |
| return config | |