|
| 1 | + |
| 2 | +""" |
| 3 | +Utility functions and helper modules |
| 4 | +
|
| 5 | +This module contains utility functions used across the framework. |
| 6 | +""" |
| 7 | + |
| 8 | +# Placeholder for future utility functions |
| 9 | +# Example utilities that could be added: |
| 10 | +# - Data preprocessing helpers |
| 11 | +# - Metric computation utilities |
| 12 | +# - Visualization helpers |
| 13 | +# - File I/O utilities |
| 14 | + |
| 15 | +__all__ = [] |
| 16 | + |
| 17 | +def get_device(local_rank=0): |
| 18 | + """Get PyTorch device for given local rank""" |
| 19 | + import torch |
| 20 | + if torch.cuda.is_available(): |
| 21 | + return torch.device(f'cuda:{local_rank}') |
| 22 | + return torch.device('cpu') |
| 23 | + |
| 24 | +def set_seed(seed=42): |
| 25 | + """Set random seed for reproducibility""" |
| 26 | + import torch |
| 27 | + import numpy as np |
| 28 | + import random |
| 29 | + |
| 30 | + random.seed(seed) |
| 31 | + np.random.seed(seed) |
| 32 | + torch.manual_seed(seed) |
| 33 | + if torch.cuda.is_available(): |
| 34 | + torch.cuda.manual_seed_all(seed) |
| 35 | + |
| 36 | +def count_parameters(model): |
| 37 | + """Count trainable parameters in model""" |
| 38 | + return sum(p.numel() for p in model.parameters() if p.requires_grad) |
| 39 | + |
| 40 | +def format_time(seconds): |
| 41 | + """Format seconds to human readable time""" |
| 42 | + hours = int(seconds // 3600) |
| 43 | + minutes = int((seconds % 3600) // 60) |
| 44 | + secs = int(seconds % 60) |
| 45 | + return f"{hours:02d}:{minutes:02d}:{secs:02d}" |
| 46 | + |
| 47 | +def get_gpu_memory_usage(): |
| 48 | + """Get current GPU memory usage""" |
| 49 | + import torch |
| 50 | + if torch.cuda.is_available(): |
| 51 | + allocated = torch.cuda.memory_allocated() / 1024**3 # GB |
| 52 | + reserved = torch.cuda.memory_reserved() / 1024**3 |
| 53 | + return { |
| 54 | + 'allocated_gb': allocated, |
| 55 | + 'reserved_gb': reserved, |
| 56 | + 'free_gb': torch.cuda.get_device_properties(0).total_memory / 1024**3 - allocated |
| 57 | + } |
| 58 | + return None |
| 59 | + |
| 60 | +__all__.extend([ |
| 61 | + 'get_device', |
| 62 | + 'set_seed', |
| 63 | + 'count_parameters', |
| 64 | + 'format_time', |
| 65 | + 'get_gpu_memory_usage', |
| 66 | +]) |
0 commit comments