| """Shared training helpers.""" |
|
|
| from __future__ import annotations |
|
|
| import logging |
| import os |
| import random |
| import sys |
| from datetime import datetime |
|
|
| import torch |
| import torch.nn as nn |
| import yaml |
|
|
|
|
| def setup_logger(log_path: str): |
| log_formatter = logging.Formatter("%(asctime)s [%(levelname)-5.5s] %(message)s") |
| root = logging.getLogger() |
| root.setLevel(logging.INFO) |
| root.handlers.clear() |
| fh = logging.FileHandler(log_path, encoding="utf-8") |
| fh.setFormatter(log_formatter) |
| root.addHandler(fh) |
| sh = logging.StreamHandler(sys.stdout) |
| sh.setFormatter(log_formatter) |
| root.addHandler(sh) |
|
|
|
|
| def set_seed(seed: int): |
| random.seed(seed) |
| torch.manual_seed(seed) |
| torch.cuda.manual_seed_all(seed) |
|
|
|
|
| class AverageMeter: |
| def __init__(self): |
| self.reset() |
|
|
| def reset(self): |
| self.val = 0 |
| self.avg = 0 |
| self.sum = 0 |
| self.count = 0 |
|
|
| def update(self, val, n=1): |
| self.val = val |
| self.sum += val * n |
| self.count += n |
| self.avg = self.sum / max(self.count, 1) |
|
|
|
|
| class CustomDataParallel(nn.DataParallel): |
| def __getattr__(self, key): |
| try: |
| return super().__getattr__(key) |
| except AttributeError: |
| return getattr(self.module, key) |
|
|
|
|
| def load_yaml_config(path: str) -> dict: |
| with open(path) as f: |
| return yaml.safe_load(f) |
|
|
|
|
| def save_checkpoint(state, is_best, out_dir, filename="checkpoint.pth.tar"): |
| os.makedirs(out_dir, exist_ok=True) |
| path = os.path.join(out_dir, filename) |
| torch.save(state, path) |
| if is_best: |
| best = os.path.join(out_dir, "checkpoint_best_loss.pth.tar") |
| torch.save(state, best) |
| logging.info(f"Saved checkpoint to {path} (best={is_best})") |
|
|
|
|
| def load_checkpoint_dict(path: str, map_location="cpu"): |
| ckpt = torch.load(path, map_location=map_location) |
| if isinstance(ckpt, dict) and "state_dict" in ckpt: |
| state = ckpt["state_dict"] |
| else: |
| state = ckpt |
| |
| out = {} |
| for k, v in state.items(): |
| out[k[7:] if k.startswith("module.") else k] = v |
| return out, ckpt if isinstance(ckpt, dict) else {"state_dict": state} |
|
|
|
|
| def adamw_trainable(model: nn.Module, lr: float, weight_decay: float = 0.01): |
| params = [p for p in model.parameters() if p.requires_grad] |
| return torch.optim.AdamW(params, lr=lr, weight_decay=weight_decay) |
|
|
|
|
| def exp_dir(root: str, exp_name: str, quality_level) -> str: |
| path = os.path.join(root, exp_name, str(quality_level)) |
| os.makedirs(path, exist_ok=True) |
| return path |
|
|
