from __future__ import annotations import math import os import torch import torch.optim as optim from flexibrain.config import RunConfig from flexibrain.data import build_pretrain_dataloaders from flexibrain.data.collate import prepare_batch_data from flexibrain.distributed import cleanup_distributed, setup_distributed from flexibrain.models import build_pretrain_model from flexibrain.utils.logging import setup_logger from flexibrain.utils.seed import set_seed from flexibrain.utils.training import get_dynamic_momentum, update_ema class Pretrainer: def __init__(self, cfg: RunConfig): self.cfg = cfg self.rank = cfg.training.local_rank self.world_size = cfg.training.world_size if self.world_size > 1: setup_distributed(self.rank, self.world_size) self.device = torch.device(f"cuda:{self.rank}" if torch.cuda.is_available() else "cpu") self.logger = setup_logger("pretrain", cfg.logging.log_dir, rank=self.rank) def build(self): set_seed(self.cfg.training.seed) self.model = build_pretrain_model(self.cfg.model, self.device) self.train_loader, self.val_loader = build_pretrain_dataloaders(self.cfg.data, self.cfg.training, rank=self.rank, world_size=self.world_size) self.optimizer = optim.AdamW(self.model.parameters(), lr=self.cfg.training.lr, weight_decay=self.cfg.training.weight_decay) self.use_amp = bool(self.cfg.training.use_amp and self.device.type == "cuda") self.scaler = torch.cuda.amp.GradScaler(enabled=self.use_amp) total_steps = max(1, len(self.train_loader) * self.cfg.training.epochs) warmup_steps = max(1, len(self.train_loader) * self.cfg.training.warmup_epochs) def lr_lambda(step): if step < warmup_steps: return step / warmup_steps progress = (step - warmup_steps) / max(1, total_steps - warmup_steps) cycle_progress = (progress * 4) % 1.0 if cycle_progress < 0.8: return 0.5 * (1 + math.cos(math.pi * cycle_progress / 0.8)) return 0.0 self.scheduler = optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda) return self def _optimizer_step(self, momentum: float) -> None: if self.cfg.training.grad_clip > 0: self.scaler.unscale_(self.optimizer) torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.cfg.training.grad_clip) self.scaler.step(self.optimizer) self.scaler.update() self.optimizer.zero_grad(set_to_none=True) update_ema(self.model, momentum) self.scheduler.step() def train_one_epoch(self, epoch: int) -> float: self.model.train() total_loss = 0.0 num_batches = 0 accumulation_steps = self.cfg.training.grad_accumulation_steps momentum = get_dynamic_momentum(epoch, self.cfg.training.epochs, self.cfg.model.momentum, self.cfg.model.final_momentum) self.optimizer.zero_grad(set_to_none=True) for batch_idx, batch in enumerate(self.train_loader): x, meta, orig_Ts, affines = prepare_batch_data(batch, self.device) with torch.autocast(device_type=self.device.type, dtype=torch.float16, enabled=self.use_amp): loss, _, _, _ = self.model(x, mask_ratio=self.cfg.training.mask_ratio, meta=meta, orig_Ts=orig_Ts, affines=affines) self.scaler.scale(loss / accumulation_steps).backward() total_loss += float(loss.item()) num_batches += 1 if (batch_idx + 1) % accumulation_steps == 0: self._optimizer_step(momentum) if self.rank == 0 and (batch_idx + 1) % self.cfg.logging.log_interval == 0: self.logger.info("Epoch %d [%d/%d] loss=%.6f avg=%.6f momentum=%.6f", epoch + 1, batch_idx + 1, len(self.train_loader), loss.item(), total_loss / num_batches, momentum) if num_batches % accumulation_steps != 0: self._optimizer_step(momentum) return total_loss / max(1, num_batches) @torch.no_grad() def validate(self, epoch: int) -> float: self.model.eval() total_loss = 0.0 num_batches = 0 for batch in self.val_loader: x, meta, orig_Ts, affines = prepare_batch_data(batch, self.device) with torch.autocast(device_type=self.device.type, dtype=torch.float16, enabled=self.use_amp): loss, _, _, _ = self.model(x, mask_ratio=self.cfg.training.mask_ratio, meta=meta, orig_Ts=orig_Ts, affines=affines) total_loss += float(loss.item()) num_batches += 1 avg = total_loss / max(1, num_batches) if self.rank == 0: self.logger.info("Epoch %d validation loss=%.6f", epoch + 1, avg) return avg def save(self, epoch: int, val_loss: float, best_loss: float, is_best: bool): if self.rank != 0: return os.makedirs(self.cfg.logging.checkpoint_dir, exist_ok=True) payload = { "epoch": epoch, "model_state_dict": self.model.state_dict(), "optimizer_state_dict": self.optimizer.state_dict(), "scheduler_state_dict": self.scheduler.state_dict(), "val_loss": val_loss, "best_loss": best_loss, "config": vars(self.cfg.model), } torch.save(payload, os.path.join(self.cfg.logging.checkpoint_dir, "checkpoint_latest.pt")) if is_best: torch.save(payload, os.path.join(self.cfg.logging.checkpoint_dir, "checkpoint_best.pt")) def fit(self): self.build() if self.rank == 0: self.logger.info("Starting pretrain on %s", self.device) self.logger.info("Train size=%d Val size=%d", len(self.train_loader.dataset), len(self.val_loader.dataset)) best_loss = float("inf") for epoch in range(self.cfg.training.epochs): if hasattr(self.train_loader.sampler, "set_epoch"): self.train_loader.sampler.set_epoch(epoch) train_loss = self.train_one_epoch(epoch) val_loss = self.validate(epoch) is_best = val_loss < best_loss if is_best: best_loss = val_loss self.save(epoch, val_loss, best_loss, is_best=is_best) if self.rank == 0: self.logger.info("Epoch %d done train=%.6f val=%.6f best=%.6f", epoch + 1, train_loss, val_loss, best_loss) if self.world_size > 1: cleanup_distributed()