| """I/O utilities for checkpoint management.""" |
|
|
| import os |
| import torch |
| import torch.nn as nn |
| from typing import Dict, Optional |
|
|
|
|
| def save_checkpoint(model: nn.Module, optimizer, scheduler, epoch: int, |
| losses: Dict, output_path: str, config: Dict = None): |
| """Save training checkpoint.""" |
| os.makedirs(os.path.dirname(output_path) or '.', exist_ok=True) |
|
|
| state = { |
| 'epoch': epoch, |
| 'model_state_dict': model.state_dict(), |
| 'optimizer_state_dict': optimizer.state_dict(), |
| 'scheduler_state_dict': scheduler.state_dict(), |
| 'losses': losses, |
| } |
| if config: |
| state['config'] = config |
|
|
| torch.save(state, output_path) |
| return output_path |
|
|
|
|
| def load_checkpoint(model: nn.Module, checkpoint_path: str, |
| optimizer=None, scheduler=None, |
| strict: bool = True) -> Dict: |
| """Load checkpoint and optionally restore optimizer/scheduler state.""" |
| checkpoint = torch.load(checkpoint_path, map_location='cpu') |
|
|
| state_dict = checkpoint.get('model_state_dict', checkpoint) |
| model.load_state_dict(state_dict, strict=strict) |
|
|
| if optimizer and 'optimizer_state_dict' in checkpoint: |
| optimizer.load_state_dict(checkpoint['optimizer_state_dict']) |
|
|
| if scheduler and 'scheduler_state_dict' in checkpoint: |
| scheduler.load_state_dict(checkpoint['scheduler_state_dict']) |
|
|
| return checkpoint |
|
|
