| """Training loop: exactly ``training.max_steps`` optimizer updates, with EMA weights, gradient clipping and |
| validation events. |
| |
| A validation event fires after global update ``step`` when ``step % val_every_updates == 0``, at ``max_steps`` and |
| at every step listed in ``val_extra_steps`` (0 = the untrained model). An event evaluates the training objective on |
| the validation pairs (``val_loss``, and ``val_loss_ema`` for the EMA weights), writes one row of ``train_log.csv`` |
| (``event`` = 1), saves ``checkpoints/checkpoint_step<NNNNNNNN>.pt`` and the full-state ``restart_latest.pt``. With |
| ``log_loss_every_epochs: k`` a loss-only row (``event`` = 0, no checkpoint) is written at the end of every k-th epoch. |
| ``train_loss`` of a row is the mean over the updates since the previous row. The validation passes run under a |
| forked RNG with a fixed seed and restore the loader's generator state, so they never perturb the training |
| trajectory. ``checkpoints/checkpoint.pt`` holds the final weights; every checkpoint stores the raw and the EMA |
| weights and the run configuration. |
| |
| EMA and LR ramp-up follow the convention of NVlabs/edm (``training/training_loop.py``): ``ema_halflife_kimg``, |
| ``ema_rampup_ratio`` and ``lr_rampup_kimg`` come from the training section (overrides) or the method section. |
| """ |
| from __future__ import annotations |
|
|
| import copy |
| import os |
| import random |
| import time |
| from pathlib import Path |
| from typing import Any, Dict |
|
|
| import numpy as np |
| import torch |
|
|
| RESTART_SCHEMA = "restart_v1" |
| RESTART_NAME = "restart_latest.pt" |
| LOG_COLUMNS = "step,epoch,batch_index,event,train_loss,val_loss,val_loss_ema,train_eval_loss,train_eval_loss_ema,lr,elapsed_s" |
|
|
|
|
| def _ema_beta(cur_nimg: int, batch_size: int, ema_halflife_kimg: float, ema_rampup_ratio: float | None) -> float: |
| ema_halflife_nimg = ema_halflife_kimg * 1000 |
| if ema_rampup_ratio is not None: |
| ema_halflife_nimg = min(ema_halflife_nimg, cur_nimg * ema_rampup_ratio) |
| return 0.5 ** (batch_size / max(ema_halflife_nimg, 1e-8)) |
|
|
|
|
| def _set_rampup_lr(optimizer: torch.optim.Optimizer, base_lr: float, cur_nimg: int, lr_rampup_kimg: float) -> None: |
| factor = min(cur_nimg / max(lr_rampup_kimg * 1000, 1e-8), 1.0) |
| for g in optimizer.param_groups: |
| g["lr"] = base_lr * factor |
|
|
|
|
| def _eval_loss(model, loader, method, device, max_batches=None) -> float: |
| losses = [] |
| with torch.no_grad(): |
| for i, batch in enumerate(loader): |
| if max_batches is not None and i >= max_batches: |
| break |
| batch = tuple(t.to(device) for t in batch) |
| losses.append(method.loss(model, batch).item()) |
| return float(np.mean(losses)) if losses else float("nan") |
|
|
|
|
| def _save_checkpoint(path: Path, model, ema_model, cfg: Dict[str, Any], epoch: int, step: int) -> None: |
| payload = {"model_state_dict": model.state_dict(), "epoch": epoch, "config": cfg, "step": int(step)} |
| if ema_model is not None: |
| payload["ema_state_dict"] = ema_model.state_dict() |
| torch.save(payload, path) |
|
|
|
|
| def _save_restart(path: Path, payload: Dict[str, Any]) -> None: |
| tmp = path.with_name(path.name + ".tmp") |
| torch.save(payload, tmp) |
| os.replace(tmp, path) |
|
|
|
|
| def _truncate_log(log_path: Path, keep_step: int) -> int: |
| """Drop the rows after ``keep_step`` (and a partial final line).""" |
| lines = log_path.read_text().splitlines() |
| header, rows = lines[0], lines[1:] |
| kept = [] |
| for line in rows: |
| try: |
| if int(line.split(",", 1)[0]) <= keep_step: |
| kept.append(line) |
| except ValueError: |
| continue |
| log_path.write_text("\n".join([header] + kept) + "\n") |
| return len(rows) - len(kept) |
|
|
|
|
| class Trainer: |
| def __init__(self, model, method, train_loader, val_loader, cfg: Dict[str, Any], run_dir: Path, device: str): |
| self.model = model |
| self.method = method |
| self.train_loader = train_loader |
| self.val_loader = val_loader |
| self.cfg = cfg |
| self.run_dir = Path(run_dir) |
| self.device = device |
|
|
| tr, m = cfg["training"], cfg["method"] |
| allow_tf32 = tr.get("allow_tf32") |
| if allow_tf32 is not None: |
| torch.backends.cuda.matmul.allow_tf32 = bool(allow_tf32) |
| torch.backends.cudnn.allow_tf32 = bool(allow_tf32) |
| print(f"precision: allow_tf32={bool(allow_tf32)} (matmul + cudnn)") |
| if bool(tr.get("deterministic", False)): |
| os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8") |
| torch.backends.cudnn.deterministic = True |
| torch.backends.cudnn.benchmark = False |
| torch.use_deterministic_algorithms(True) |
| print("determinism: use_deterministic_algorithms(True), cudnn.deterministic, benchmark off, " |
| "CUBLAS_WORKSPACE_CONFIG=" + os.environ["CUBLAS_WORKSPACE_CONFIG"]) |
|
|
| self.batch_size = int(tr["batch_size"]) |
| if "lr" in tr: |
| self.base_lr = float(tr["lr"]) |
| elif "lr" in m: |
| self.base_lr = float(m["lr"]) |
| else: |
| raise KeyError("'lr' must be set in the training or the method section") |
| self.val_seed = int(tr.get("val_seed", 777)) |
| self.restart_checkpointing = bool(tr.get("restart_checkpointing", True)) |
|
|
| if tr.get("max_steps") is None or tr.get("val_every_updates") is None: |
| raise KeyError("training.max_steps and training.val_every_updates are required") |
| self.max_steps = int(tr["max_steps"]) |
| self.val_every_updates = int(tr["val_every_updates"]) |
| if self.max_steps <= 0 or self.val_every_updates <= 0: |
| raise ValueError("max_steps and val_every_updates must be positive") |
| self.val_extra_steps = tr.get("val_extra_steps") |
| if self.val_extra_steps is not None: |
| if not isinstance(self.val_extra_steps, (list, tuple)) or \ |
| any(isinstance(s, bool) or int(s) != s for s in self.val_extra_steps): |
| raise ValueError("training.val_extra_steps must be a list of integers") |
| xs = sorted({int(s) for s in self.val_extra_steps}) |
| bad = [s for s in xs if s < 0 or s >= self.max_steps] |
| if bad: |
| raise ValueError(f"training.val_extra_steps must lie in [0, max_steps): {bad}") |
| self.val_extra_steps = xs |
| self.log_loss_every_epochs = tr.get("log_loss_every_epochs") |
| if self.log_loss_every_epochs is not None: |
| k = self.log_loss_every_epochs |
| if isinstance(k, bool) or int(k) != k or int(k) <= 0: |
| raise ValueError("training.log_loss_every_epochs must be a positive integer") |
| self.log_loss_every_epochs = int(k) |
|
|
| def _opt(key, default=None): |
| return tr[key] if key in tr else m.get(key, default) |
|
|
| self.ema_halflife_kimg = _opt("ema_halflife_kimg") |
| self.ema_rampup_ratio = _opt("ema_rampup_ratio", 0.05) |
| self.lr_rampup_kimg = _opt("lr_rampup_kimg") |
| self.max_grad_norm = _opt("max_grad_norm") |
| self.weight_decay = _opt("weight_decay") |
| self.use_ema = self.ema_halflife_kimg is not None and float(self.ema_halflife_kimg) > 0 |
| self.use_lr_rampup = self.lr_rampup_kimg is not None and float(self.lr_rampup_kimg) > 0 |
|
|
| def fit(self, resume: Dict[str, Any] | None = None) -> None: |
| self.run_dir.mkdir(parents=True, exist_ok=True) |
| ckpt_dir = self.run_dir / "checkpoints" |
| ckpt_dir.mkdir(parents=True, exist_ok=True) |
|
|
| model = self.model.to(self.device) |
| if self.use_ema: |
| ema_model = copy.deepcopy(model).eval().to(self.device) |
| for p in ema_model.parameters(): |
| p.requires_grad_(False) |
| else: |
| ema_model = None |
| optimizer = self.method.configure_optimizer(model, {"lr": self.base_lr, "weight_decay": self.weight_decay}) |
|
|
| start_epoch = 1 |
| if resume is not None: |
| if resume.get("schema") != RESTART_SCHEMA: |
| raise AssertionError(f"restart checkpoint schema {resume.get('schema')!r} != {RESTART_SCHEMA}") |
| model.load_state_dict(resume["model_state_dict"]) |
| if ema_model is not None: |
| if resume.get("ema_state_dict") is None: |
| raise AssertionError("EMA is on but the restart checkpoint has no EMA state") |
| ema_model.load_state_dict(resume["ema_state_dict"]) |
| optimizer.load_state_dict(resume["optimizer_state_dict"]) |
| start_epoch = int(resume["epoch"]) |
|
|
| n_train_pairs = len(self.train_loader.dataset) |
| steps_per_epoch = len(self.train_loader) |
| n_epochs = max(1, -(-self.max_steps // steps_per_epoch)) |
| print(f"schedule: max_steps = {self.max_steps}, validation + checkpoint every {self.val_every_updates} updates " |
| f"(+ at max_steps); {n_train_pairs} training pairs, batch {self.batch_size}, {steps_per_epoch} steps/epoch -> " |
| f"{n_epochs} epochs (last one partial: {self.max_steps - (n_epochs - 1) * steps_per_epoch} steps)") |
| if self.use_ema: |
| seen_kimg = (self.max_steps - 1) * self.batch_size / 1000.0 |
| eff_kimg = min(float(self.ema_halflife_kimg), float(self.ema_rampup_ratio) * seen_kimg) |
| print(f"EMA: half-life {self.ema_halflife_kimg} kimg, ramp-up ratio {self.ema_rampup_ratio} " |
| f"(effective half-life at the last update ~= {eff_kimg:.1f} kimg)") |
| else: |
| print("EMA: disabled") |
| print(f"LR ramp-up: {self.lr_rampup_kimg if self.use_lr_rampup else 'disabled'}; " |
| f"gradient clipping: {self.max_grad_norm if self.max_grad_norm is not None else 'disabled'}") |
|
|
| log_path = self.run_dir / "train_log.csv" |
| if resume is not None: |
| if int(resume["step"]) >= self.max_steps: |
| raise AssertionError(f"resume step {int(resume['step'])} >= max_steps {self.max_steps} (nothing to continue)") |
| if not log_path.exists(): |
| raise FileNotFoundError(f"--resume requires the existing {log_path}") |
| dropped = _truncate_log(log_path, int(resume["step"])) |
| print(f"resume: continuing at step {int(resume['step'])} (epoch {start_epoch}, batch " |
| f"{int(resume['batch_index'])}; {dropped} log row(s) dropped)") |
| else: |
| log_path.write_text(LOG_COLUMNS + "\n") |
|
|
| cur_nimg = 0 |
| if resume is not None: |
| cur_nimg = int(resume["cur_nimg"]) |
| torch.set_rng_state(resume["torch_rng_state"]) |
| if str(self.device).startswith("cuda"): |
| torch.cuda.set_rng_state_all(resume["cuda_rng_state_all"]) |
| np.random.set_state(resume["numpy_rng_state"]) |
| random.setstate(resume["python_rng_state"]) |
| gen = self.train_loader.generator |
| if resume.get("loader_generator_state_epoch_start") is not None: |
| if gen is None: |
| raise AssertionError("restart checkpoint has a loader generator state but the train loader has no generator") |
| gen.set_state(resume["loader_generator_state_epoch_start"]) |
| print(f"resume: restored optimizer, EMA, counters (cur_nimg={cur_nimg}), RNG streams and loader generator", flush=True) |
| self._fit_updates(model=model, ema_model=ema_model, optimizer=optimizer, ckpt_dir=ckpt_dir, log_path=log_path, |
| resume=resume, start_epoch=start_epoch, cur_nimg=cur_nimg, n_val_batches=len(self.val_loader)) |
|
|
| def _fit_updates(self, *, model, ema_model, optimizer, ckpt_dir, log_path, resume, start_epoch, cur_nimg, n_val_batches): |
| max_steps, every = self.max_steps, self.val_every_updates |
| extra_steps = set(self.val_extra_steps or []) |
| loss_every = int(self.log_loss_every_epochs) if self.log_loss_every_epochs else 0 |
| devices_all = [self.device] if str(self.device).startswith("cuda") else [] |
| losses_since_epoch_start: list[float] = [] |
| last_event_step = -1 |
| t_row = time.time() |
| if resume is not None: |
| losses_since_epoch_start = [float(x) for x in resume.get("losses_since_epoch_start", [])] |
| last_event_step = int(resume["step"]) |
| step, start_batch = 0, 0 |
| if resume is not None: |
| step, start_batch = int(resume["step"]), int(resume["batch_index"]) |
| epoch = start_epoch |
| losses_since_event: list[float] = [] |
| t_event = time.time() |
| done = False |
| gen = self.train_loader.generator |
|
|
| def _event(epoch_i: int, batch_i: int, gen_state_epoch_start): |
| nonlocal losses_since_event, t_event, last_event_step, t_row |
| train_loss = float(np.mean(losses_since_event)) if losses_since_event else float("nan") |
| cur_lr = optimizer.param_groups[0]["lr"] |
| model.eval() |
| g_state = gen.get_state() if gen is not None else None |
| with torch.random.fork_rng(devices=devices_all): |
| torch.manual_seed(self.val_seed) |
| val_loss = _eval_loss(model, self.val_loader, self.method, self.device) |
| train_eval = _eval_loss(model, self.train_loader, self.method, self.device, max_batches=n_val_batches) |
| if self.use_ema: |
| val_loss_ema = _eval_loss(ema_model, self.val_loader, self.method, self.device) |
| train_eval_ema = _eval_loss(ema_model, self.train_loader, self.method, self.device, max_batches=n_val_batches) |
| else: |
| val_loss_ema = train_eval_ema = float("nan") |
| if g_state is not None: |
| gen.set_state(g_state) |
| elapsed = time.time() - t_event |
| t_event = t_row = time.time() |
| last_event_step = step |
| print(f"step {step} (epoch {epoch_i}, batch {batch_i}): train_loss={train_loss:.6f} val={val_loss:.6f} " |
| f"val_ema={val_loss_ema:.6f} lr={cur_lr:.2e} elapsed={elapsed:.2f}s", flush=True) |
| with open(log_path, "a") as f: |
| f.write(f"{step},{epoch_i},{batch_i},1,{train_loss},{val_loss},{val_loss_ema}," |
| f"{train_eval},{train_eval_ema},{cur_lr},{elapsed}\n") |
| losses_since_event = [] |
| ckpt = ckpt_dir / f"checkpoint_step{step:08d}.pt" |
| _save_checkpoint(ckpt, model, ema_model, self.cfg, epoch_i, step=step) |
| print(f" saved: {ckpt.name}", flush=True) |
| if self.restart_checkpointing: |
| _save_restart(ckpt_dir / RESTART_NAME, { |
| "schema": RESTART_SCHEMA, "epoch": epoch_i, "step": step, "batch_index": batch_i, "cur_nimg": cur_nimg, |
| "model_state_dict": model.state_dict(), |
| "ema_state_dict": (ema_model.state_dict() if ema_model is not None else None), |
| "optimizer_state_dict": optimizer.state_dict(), |
| "torch_rng_state": torch.get_rng_state(), |
| "cuda_rng_state_all": (torch.cuda.get_rng_state_all() if str(self.device).startswith("cuda") else []), |
| "numpy_rng_state": np.random.get_state(), "python_rng_state": random.getstate(), |
| "loader_generator_state_epoch_start": gen_state_epoch_start, |
| "losses_since_epoch_start": list(losses_since_epoch_start), |
| "config": self.cfg}) |
| model.train() |
|
|
| def _loss_row(epoch_i: int, batch_i: int): |
| nonlocal losses_since_epoch_start, t_row |
| if not losses_since_epoch_start or step == last_event_step: |
| losses_since_epoch_start = [] |
| return |
| train_loss = float(np.mean(losses_since_epoch_start)) |
| cur_lr = optimizer.param_groups[0]["lr"] |
| model.eval() |
| g_state = gen.get_state() if gen is not None else None |
| with torch.random.fork_rng(devices=devices_all): |
| torch.manual_seed(self.val_seed) |
| val_loss = _eval_loss(model, self.val_loader, self.method, self.device) |
| val_loss_ema = (_eval_loss(ema_model, self.val_loader, self.method, self.device) if self.use_ema else float("nan")) |
| if g_state is not None: |
| gen.set_state(g_state) |
| model.train() |
| nan = float("nan") |
| elapsed = time.time() - t_row |
| t_row = time.time() |
| with open(log_path, "a") as f: |
| f.write(f"{step},{epoch_i},{batch_i},0,{train_loss},{val_loss},{val_loss_ema},{nan},{nan},{cur_lr},{elapsed}\n") |
| losses_since_epoch_start = [] |
|
|
| if 0 in extra_steps and resume is None and step == 0: |
| |
| g_state = gen.get_state() if gen is not None else None |
| model.eval() |
| with torch.random.fork_rng(devices=devices_all): |
| torch.manual_seed(self.val_seed) |
| with torch.no_grad(): |
| b0 = tuple(t.to(self.device) for t in next(iter(self.train_loader))) |
| loss0 = float(self.method.loss(model, b0).item()) |
| model.train() |
| if g_state is not None: |
| gen.set_state(g_state) |
| losses_since_event = [loss0] |
| _event(epoch, 0, gen.get_state() if gen is not None else None) |
|
|
| while not done: |
| gen_state_epoch_start = gen.get_state() if gen is not None else None |
| replay = resume is not None and epoch == start_epoch and start_batch > 0 |
| if replay: |
| gen_state_epoch_start = resume["loader_generator_state_epoch_start"] |
| model.train() |
| if not replay: |
| losses_since_epoch_start = [] |
| it = iter(self.train_loader) |
| skip = start_batch if (resume is not None and epoch == start_epoch) else 0 |
| for _ in range(skip): |
| next(it) |
| batch_i = skip |
| for batch in it: |
| batch = tuple(t.to(self.device) for t in batch) |
| if self.use_lr_rampup: |
| _set_rampup_lr(optimizer, self.base_lr, cur_nimg, float(self.lr_rampup_kimg)) |
| optimizer.zero_grad() |
| loss = self.method.loss(model, batch) |
| loss.backward() |
| if self.max_grad_norm is not None: |
| torch.nn.utils.clip_grad_norm_(model.parameters(), float(self.max_grad_norm)) |
| optimizer.step() |
| if self.use_ema: |
| beta = _ema_beta(cur_nimg, self.batch_size, float(self.ema_halflife_kimg), self.ema_rampup_ratio) |
| for p_ema, p_net in zip(ema_model.parameters(), model.parameters()): |
| p_ema.copy_(p_net.detach().lerp(p_ema, beta)) |
| for b_ema, b_net in zip(ema_model.buffers(), model.buffers()): |
| b_ema.copy_(b_net) |
| cur_nimg += batch[0].shape[0] |
| _li = loss.item() |
| losses_since_event.append(_li) |
| losses_since_epoch_start.append(_li) |
| step += 1 |
| batch_i += 1 |
| if step % every == 0 or step == max_steps or step in extra_steps: |
| _event(epoch, batch_i, gen_state_epoch_start) |
| if step >= max_steps: |
| done = True |
| break |
| if not done and loss_every and epoch % loss_every == 0: |
| _loss_row(epoch, batch_i) |
| epoch += 1 |
|
|
| print(f"trained exactly {max_steps} updates", flush=True) |
| _save_checkpoint(ckpt_dir / "checkpoint.pt", model, ema_model, self.cfg, epoch - 1, step=max_steps) |
|
|
