Coverage for widitapp/training.py: 99.65%
284 statements
« prev ^ index » next coverage.py v7.10.6, created at 2026-09-22 07:22 +0000
1from collections import OrderedDict
2import logging
3import math
4import os
5from glob import glob
6from copy import deepcopy
7from time import time
8from typing import Optional, Dict
9import torch
10import torch.nn as nn
11from torch.utils.data import DataLoader
12from accelerate import Accelerator
13from rich.progress import track
14from widit import WiDiT
15try:
16 import wandb # optional
17except Exception:
18 wandb = None
20from .diffusion import create_diffusion
23def create_logger(logging_dir: str):
24 logger = logging.getLogger("train")
25 logger.setLevel(logging.INFO)
26 # Prevent duplicate handlers
27 if not logger.handlers:
28 fmt = logging.Formatter("[%(asctime)s] %(message)s", "%Y-%m-%d %H:%M:%S")
29 sh = logging.StreamHandler()
30 sh.setFormatter(fmt)
31 fh = logging.FileHandler(f"{logging_dir}/log.txt")
32 fh.setFormatter(fmt)
33 logger.addHandler(sh)
34 logger.addHandler(fh)
35 return logger
38def get_state_dict_for_saving(model):
39 return model.module.state_dict() if hasattr(model, "module") else model.state_dict()
42def build_loss_fn(loss_fn: str | torch.nn.Module):
43 if isinstance(loss_fn, torch.nn.Module):
44 return loss_fn
45 name = loss_fn.lower()
46 if name in {"mse", "l2"}:
47 return torch.nn.MSELoss(reduction="mean")
48 if name in {"smoothl1", "huber"}:
49 return torch.nn.SmoothL1Loss(reduction="mean")
50 raise ValueError(f"Unknown loss_fn '{loss_fn}'. Expected 'mse' or 'smoothl1'.")
53def build_val_log_payload(
54 val_loss: float | Dict[str, float],
55 *,
56 use_diffusion: bool,
57 epoch: int,
58 train_steps: int,
59) -> tuple[float, Dict[str, float]]:
60 val_loss_value = val_loss.get("loss", float("nan")) if isinstance(val_loss, dict) else val_loss
61 payload = {"val/loss": val_loss_value, "epoch": epoch, "global_step": train_steps}
62 if isinstance(val_loss, dict):
63 payload.update({f"val/{name}": value for name, value in val_loss.items()})
64 return val_loss_value, payload
67def is_loss_finite(
68 loss: torch.Tensor,
69 accelerator: Accelerator,
70) -> bool:
71 local_is_finite = torch.isfinite(loss.detach()).all().to(device=accelerator.device, dtype=torch.float32)
72 return accelerator.reduce(local_is_finite, reduction="mean").item() == 1.0
75def loss_value_for_logging(loss: torch.Tensor) -> float:
76 return loss.detach().item() if loss.detach().numel() == 1 else float("nan")
79def is_finite_tensor(
80 tensor: torch.Tensor,
81 accelerator: Accelerator,
82) -> bool:
83 local_is_finite = torch.isfinite(tensor.detach()).all().to(device=accelerator.device, dtype=torch.float32)
84 return accelerator.reduce(local_is_finite, reduction="mean").item() == 1.0
87def is_cuda_out_of_memory(error: BaseException) -> bool:
88 out_of_memory_error = getattr(torch, "OutOfMemoryError", None)
89 if out_of_memory_error is not None and isinstance(error, out_of_memory_error):
90 return True
91 message = str(error).lower()
92 return isinstance(error, RuntimeError) and "cuda" in message and "out of memory" in message
95def tensor_info_for_logging(tensor: torch.Tensor) -> str:
96 return f"shape={tuple(tensor.shape)}, dtype={tensor.dtype}, device={tensor.device}"
99def cuda_is_available() -> bool:
100 # Indirection so tests can simulate a GPU host without patching torch.cuda
101 # globally, which makes torch internals attempt real CUDA initialisation.
102 return torch.cuda.is_available()
105def clear_after_cuda_oom(optimizer: torch.optim.Optimizer):
106 optimizer.zero_grad(set_to_none=True)
107 if cuda_is_available():
108 torch.cuda.empty_cache()
111@torch.no_grad()
112def update_ema(ema_model, model, decay=0.9999):
113 ema_params = OrderedDict(ema_model.named_parameters())
114 model_params = OrderedDict(model.named_parameters())
115 for name, param in model_params.items():
116 name = name.replace("module.", "")
117 ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay)
120def requires_grad(model, flag=True):
121 for p in model.parameters():
122 p.requires_grad = flag
125@torch.no_grad()
126def _run_validation_loop(
127 accelerator: Accelerator,
128 model_for_eval: torch.nn.Module,
129 diffusion,
130 dataloader: DataLoader,
131 device: torch.device,
132 dtype: torch.dtype,
133 use_diffusion: bool,
134 criterion: torch.nn.Module,
135 extra_criteria: Dict[str, torch.nn.Module] | None = None,
136 timestep_seed: int = 42,
137 diffusion_loss_fn: torch.nn.Module | None = None,
138) -> Dict[str, float]:
139 model_for_eval.eval()
140 for name, metric in (extra_criteria or {}).items():
141 if not isinstance(name, str) or not name or name == "loss" or name.startswith("val/"):
142 raise ValueError("Metric names must be nonempty, omit 'val/', and cannot be 'loss'.")
143 metric.to(device).eval()
144 total_loss, total_batches = 0.0, 0
145 extra_totals = {name: 0.0 for name in (extra_criteria or {})}
146 rng = torch.Generator(device=device).manual_seed(timestep_seed)
148 for batch in track(dataloader, total=len(dataloader), description="Validation:"):
149 if not isinstance(batch, (list, tuple)) or len(batch) not in {2, 3}:
150 raise ValueError("Validation dataloader must return (x, target) or (x, target, timestep).")
151 x, target = batch[0], batch[1]
152 timestep = batch[2] - 1 if len(batch) == 3 else None
154 x = x.to(device=device, dtype=dtype, non_blocking=True)
155 target = target.to(device=device, dtype=dtype, non_blocking=True)
157 extra_losses = {}
158 if use_diffusion:
159 # Mirror training: model(input=target, conditioned=x)
160 t = torch.randint(
161 0,
162 diffusion.num_timesteps,
163 (x.shape[0],),
164 device=device,
165 generator=rng,
166 )
167 noise = torch.randn(
168 target.shape,
169 dtype=target.dtype,
170 device=target.device,
171 generator=rng,
172 )
173 loss_dict = diffusion.training_losses(
174 model_for_eval,
175 target,
176 t,
177 dict(conditioned=x),
178 noise=noise,
179 **({"image_loss_fn": diffusion_loss_fn} if diffusion_loss_fn is not None else {}),
180 **({"return_pred_xstart": True} if extra_criteria else {}),
181 )
182 loss = loss_dict["mse" if diffusion_loss_fn is None else "loss"].mean()
183 if extra_criteria:
184 y = loss_dict["pred_xstart"]
185 else:
186 y = model_for_eval(x, timestep=timestep)
187 loss = criterion(y, target)
188 if extra_criteria:
189 for name, extra in extra_criteria.items():
190 extra_loss = extra(y, target)
191 if not isinstance(extra_loss, torch.Tensor) or extra_loss.ndim != 0:
192 raise ValueError(f"Metric '{name}' must return a scalar batch mean tensor.")
193 extra_loss = accelerator.reduce(extra_loss, reduction="mean")
194 extra_losses[name] = extra_loss
196 loss = accelerator.reduce(loss, reduction="mean")
197 if not is_loss_finite(loss, accelerator):
198 if accelerator.is_main_process:
199 logging.getLogger("train").warning(
200 "Skipping non-finite validation batch: "
201 f"loss={loss_value_for_logging(loss)}"
202 )
203 continue
205 total_loss += loss.item()
206 for name, extra_loss in extra_losses.items():
207 extra_totals[name] += extra_loss.item()
208 total_batches += 1
210 if total_batches <= 0:
211 return {name: float("nan") for name in ("loss", *extra_totals)}
213 result = {"loss": total_loss / total_batches}
214 for name, total in extra_totals.items():
215 result[name] = total / total_batches
216 return result
219def train(
220 model: WiDiT,
221 training_dataloader: DataLoader,
222 validation_dataloader: DataLoader | None = None,
223 results_dir: str = "results",
224 epochs: int = 40,
225 log_every: int = 10,
226 use_diffusion: bool = True,
227 loss_fn: str | torch.nn.Module = "mse",
228 precision: str = "fp16",
229 wandb_logging: bool = False,
230 wandb_project: str = "WiDiT",
231 run_name: str|None = None,
232 learning_rate: float = 1e-4,
233 wandb_config: Optional[Dict] = None,
234 wandb_log_artifacts: bool = False,
235 diffusion_loss_fn: torch.nn.Module | None = None,
236 metrics: Dict[str, torch.nn.Module] | None = None,
237):
238 torch.backends.cuda.matmul.allow_tf32 = True
239 torch.backends.cudnn.allow_tf32 = True
240 assert cuda_is_available(), "Training currently requires at least one GPU."
242 accelerator = Accelerator(mixed_precision=("no" if precision == "fp32" else precision))
243 device = accelerator.device
245 dtype_map = {
246 "no": torch.float32,
247 "fp32": torch.float32,
248 "fp16": torch.float16,
249 "bf16": torch.bfloat16,
250 }
251 train_dtype = dtype_map[precision]
253 # ---- experiment dirs & logger ----
254 if accelerator.is_main_process:
255 os.makedirs(results_dir, exist_ok=True)
256 experiment_index = len(glob(f"{results_dir}/*"))
257 model_string_name = model.__class__.__name__
258 run_name = run_name or f"{model_string_name}-{experiment_index:02d}"
259 experiment_dir = f"{results_dir}/{run_name}"
260 checkpoint_dir = f"{experiment_dir}/checkpoints"
261 os.makedirs(checkpoint_dir, exist_ok=True)
262 logger = create_logger(experiment_dir)
263 logger.info(f"Experiment directory: {experiment_dir}")
264 else:
265 logger = logging.getLogger("train") # no handlers; quiet worker
267 if wandb_logging and accelerator.is_main_process:
268 if wandb is None:
269 logger.warning("wandb=True but package not available. Skipping W&B init.")
270 else:
271 run_cfg = {
272 "precision": precision,
273 "use_diffusion": use_diffusion,
274 "epochs": epochs,
275 "log_every": log_every,
276 "loss_fn": loss_fn if isinstance(loss_fn, str) else loss_fn.__class__.__name__,
277 "optimizer": "AdamW",
278 "lr": learning_rate,
279 "weight_decay": 0.0,
280 "model": model.__class__.__name__,
281 "params": sum(p.numel() for p in model.parameters()),
282 }
283 if wandb_config:
284 run_cfg.update(wandb_config)
286 wandb.init(project=wandb_project, name=run_name, config=run_cfg, dir=experiment_dir)
287 # If you want gradients & params logged (can be heavy):
288 # wandb.watch(model, log="all", log_freq=log_every)
290 if accelerator.is_main_process:
291 logger.info(f"Training with precision={precision} (dtype={train_dtype}) on {device}")
293 criterion = build_loss_fn(loss_fn).to(device)
294 if diffusion_loss_fn is not None:
295 diffusion_loss_fn = diffusion_loss_fn.to(device)
296 extra_val_criteria = metrics if metrics is not None else ({} if use_diffusion else {
297 "mse": torch.nn.MSELoss(reduction="mean"),
298 "smoothl1": torch.nn.SmoothL1Loss(reduction="mean"),
299 })
301 # ---- models ----
302 model = model.to(device)
303 ema = deepcopy(model).to(device)
304 requires_grad(ema, False)
306 diffusion = create_diffusion(timestep_respacing="") if use_diffusion else None
308 if accelerator.is_main_process:
309 n_params = sum(p.numel() for p in model.parameters())
310 logger.info(f"WiDiT Parameters: {n_params:,}")
311 logger.info(f"Training mode: {'diffusion' if use_diffusion else 'supervised (no diffusion)'}")
313 # ---- optimizer ----
314 opt = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=0.0)
316 if accelerator.is_main_process:
317 logger.info(f"Training dataloader: {len(training_dataloader):,} batches/epoch")
318 if validation_dataloader is not None:
319 logger.info(f"Validation dataloader: {len(validation_dataloader):,} batches/epoch")
321 # ---- prepare with Accelerate ----
322 update_ema(ema, model, decay=0)
323 model, ema, opt, training_dataloader = accelerator.prepare(model, ema, opt, training_dataloader)
324 if validation_dataloader is not None:
325 validation_dataloader = accelerator.prepare(validation_dataloader)
327 # ---- training loop ----
328 train_steps = 0
329 log_steps = 0
330 running_loss = 0.0
331 start_time = time()
332 best_val_loss = float("inf")
333 best_checkpoint_saved = False
335 if accelerator.is_main_process:
336 logger.info(f"Training for {epochs} epoch(s)...")
338 for epoch in range(epochs):
339 if accelerator.is_main_process:
340 logger.info(f"Epoch {epoch}...")
342 model.train()
343 len_training_dataloader = len(training_dataloader)
344 for batch in track(training_dataloader, total=len(training_dataloader), description="Training:"):
345 # Expect (x, target)
346 if not isinstance(batch, (list, tuple)) or len(batch) not in {2, 3}:
347 raise ValueError("Training dataloader must return (x, target) or (x, target, timestep).")
348 x, target = batch[0], batch[1]
349 timestep = batch[2] - 1 if len(batch) == 3 else None
351 x = x.to(device=accelerator.device, dtype=train_dtype, non_blocking=True)
352 target = target.to(device=accelerator.device, dtype=train_dtype, non_blocking=True)
354 x_is_finite = is_finite_tensor(x, accelerator)
355 target_is_finite = is_finite_tensor(target, accelerator)
356 if not x_is_finite or not target_is_finite:
357 opt.zero_grad(set_to_none=True)
358 if accelerator.is_main_process:
359 logger.warning(
360 f"Skipping non-finite training batch at epoch={epoch}, step={train_steps + 1}: "
361 f"x_is_finite={x_is_finite}, target_is_finite={target_is_finite}"
362 )
363 continue
365 local_oom = torch.tensor(0.0, device=accelerator.device)
366 loss = None
367 try:
368 with accelerator.autocast():
369 if use_diffusion:
370 timestep = torch.randint(0, diffusion.num_timesteps, (x.shape[0],), device=accelerator.device)
371 # Mirror your earlier logic: model(input=target, conditioned=x)
372 loss_dict = diffusion.training_losses(
373 model, target, timestep, dict(conditioned=x),
374 **({"image_loss_fn": diffusion_loss_fn} if diffusion_loss_fn is not None else {}),
375 )
376 loss = loss_dict["loss"].mean()
377 else:
378 y = model(x, timestep=timestep)
379 loss = criterion(y, target)
380 except Exception as error:
381 if not is_cuda_out_of_memory(error):
382 raise
383 local_oom.fill_(1.0)
384 if cuda_is_available():
385 torch.cuda.empty_cache()
387 any_oom = accelerator.reduce(local_oom, reduction="sum").item() > 0
388 if any_oom:
389 clear_after_cuda_oom(opt)
390 if accelerator.is_main_process:
391 logger.warning(
392 f"Skipping CUDA out-of-memory training batch at epoch={epoch}, step={train_steps + 1}: "
393 f"x=({tensor_info_for_logging(x)}), target=({tensor_info_for_logging(target)})"
394 )
395 continue
397 if not is_loss_finite(loss, accelerator):
398 opt.zero_grad(set_to_none=True)
399 if accelerator.is_main_process:
400 logger.warning(
401 f"Skipping non-finite training loss at epoch={epoch}, step={train_steps + 1}: "
402 f"loss={loss_value_for_logging(loss)}"
403 )
404 continue
406 opt.zero_grad(set_to_none=True)
407 accelerator.backward(loss)
408 opt.step()
409 update_ema(ema, model)
411 running_loss += loss.item()
412 log_steps += 1
413 train_steps += 1
415 # ---- logging ----
416 if train_steps % log_every == 0:
417 torch.cuda.synchronize()
418 end_time = time()
419 steps_per_sec = log_steps / max(end_time - start_time, 1e-6)
420 avg_loss = running_loss / log_steps
421 avg_loss_tensor = torch.tensor(avg_loss, device=accelerator.device)
422 avg_loss_tensor = accelerator.reduce(avg_loss_tensor, reduction="mean")
423 avg_loss_val = avg_loss_tensor.item()
425 if accelerator.is_main_process:
426 logger.info(
427 f"(epoch={epoch:03d}, step={train_steps:07d}/{len_training_dataloader}, {train_steps/len_training_dataloader*100:.1f}%) "
428 f"train/loss={avg_loss_val:.4f} speed={steps_per_sec:.2f} steps/s"
429 )
430 if wandb_logging and wandb is not None:
431 wandb.log(
432 {
433 "train/loss": avg_loss_val,
434 "train/steps_per_sec": steps_per_sec,
435 "epoch": epoch,
436 "global_step": train_steps,
437 },
438 step=train_steps,
439 )
441 running_loss = 0.0
442 log_steps = 0
443 start_time = time()
446 # ---- validation (end of epoch) ----
447 if validation_dataloader is not None:
448 val_loss = _run_validation_loop(
449 accelerator=accelerator,
450 model_for_eval=ema, # evaluate EMA
451 diffusion=diffusion,
452 dataloader=validation_dataloader,
453 device=accelerator.device,
454 dtype=train_dtype,
455 use_diffusion=use_diffusion,
456 criterion=criterion,
457 extra_criteria=extra_val_criteria or None,
458 timestep_seed=0, # deterministic validation timesteps
459 diffusion_loss_fn=diffusion_loss_fn,
460 )
461 val_loss_value, log_payload = build_val_log_payload(
462 val_loss,
463 use_diffusion=use_diffusion,
464 epoch=epoch,
465 train_steps=train_steps,
466 )
467 if accelerator.is_main_process:
468 logger.info(f"(epoch={epoch:03d}) val/loss={val_loss_value:.4f}")
469 for name, value in val_loss.items():
470 if name != "loss":
471 logger.info(f"(epoch={epoch:03d}) val/{name}={value:.4f}")
472 if wandb_logging and wandb is not None:
473 wandb.log(log_payload, step=train_steps)
475 # Save best
476 val_loss_is_finite = math.isfinite(val_loss_value)
477 is_best = val_loss_is_finite and val_loss_value < best_val_loss
478 save_missing_best = not best_checkpoint_saved
479 if is_best or save_missing_best:
480 if val_loss_is_finite:
481 best_val_loss = val_loss_value
482 best_checkpoint_saved = True
483 best_path = f"{checkpoint_dir}/best.pt"
484 ema.save(best_path)
485 if is_best:
486 logger.info(f"New best checkpoint (val_loss={val_loss_value:.4f}) saved to {best_path}")
487 else:
488 logger.warning(
489 "Validation loss is non-finite and no best checkpoint exists; "
490 f"saved fallback checkpoint to {best_path}"
491 )
493 if wandb_logging and wandb is not None and wandb_log_artifacts:
494 art = wandb.Artifact(f"{model_string_name}-best", type="model")
495 art.add_file(best_path)
496 wandb.log_artifact(art)
498 # ---- final checkpoint ----
499 if accelerator.is_main_process:
500 if wandb_logging and wandb is not None:
501 wandb.summary["best/val_loss"] = best_val_loss
502 wandb.finish()