File size: 3,593 Bytes
a311e94 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | """Composable training algorithms.
Each algorithm implements a small set of hooks that the Agron training engine
calls in order::
algo.on_train_begin(ctx)
for step, batch in loader:
algo.on_step_begin(ctx, batch)
loss = algo.train_step(ctx, batch)
algo.on_step_end(ctx, loss)
Algorithms are deliberately composable: several can be active at once and they
will all see the same context.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any, Dict, Optional
import torch
@dataclass
class AlgorithmContext:
"""Mutable state shared between the engine and the active algorithms."""
model: Any = None
optimizer: Any = None
scheduler: Any = None
device: torch.device = torch.device("cpu")
step: int = 0
epoch: int = 0
max_steps: int = 1000
world_size: int = 1
extra: Dict[str, Any] = field(default_factory=dict)
def get(self, key: str, default: Any = None) -> Any:
return self.extra.get(key, default)
def set(self, key: str, value: Any) -> None:
self.extra[key] = value
class BaseAlgorithm:
"""No-op baseline every algorithm inherits from."""
name: str = "base"
category: str = "misc"
description: str = ""
def __init__(self, **kwargs):
self.config = kwargs
# -- lifecycle -------------------------------------------------------- #
def on_train_begin(self, ctx: AlgorithmContext) -> None:
pass
def on_train_end(self, ctx: AlgorithmContext) -> None:
pass
def on_epoch_begin(self, ctx: AlgorithmContext) -> None:
pass
def on_step_begin(self, ctx: AlgorithmContext, batch: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
return batch
def compute_loss(self, ctx: AlgorithmContext, batch: Dict[str, torch.Tensor], outputs: Any, loss: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
return loss
def train_step(self, ctx: AlgorithmContext, batch: Dict[str, torch.Tensor]) -> torch.Tensor:
"""Default single forward/backward pass."""
batch = self.on_step_begin(ctx, batch)
outputs = ctx.model(**batch)
loss = outputs["loss"] if isinstance(outputs, dict) else outputs.loss
loss = self.compute_loss(ctx, batch, outputs, loss)
loss = compose_loss(ctx, batch, outputs, loss, skip=self)
if loss is not None:
loss.backward()
return loss
def on_before_optimizer_step(self, ctx: AlgorithmContext) -> None:
pass
def should_step(self, ctx: AlgorithmContext) -> bool:
"""Whether the optimiser should run after this micro-step."""
return True
def on_step_end(self, ctx: AlgorithmContext, loss: Optional[torch.Tensor]) -> None:
pass
# -- introspection ---------------------------------------------------- #
def state_dict(self) -> Dict[str, Any]:
return {}
def load_state_dict(self, state: Dict[str, Any]) -> None:
pass
def compose_loss(
ctx: AlgorithmContext,
batch: Dict[str, torch.Tensor],
outputs: Any,
loss: Optional[torch.Tensor],
skip: Optional[BaseAlgorithm] = None,
) -> Optional[torch.Tensor]:
"""Apply every registered ``compute_loss`` hook in order.
``skip`` prevents the algorithm that already computed the loss (the step
driver, e.g. SAM or Mixup) from being applied twice.
"""
for algo in ctx.get("loss_hooks") or []:
if algo is skip:
continue
loss = algo.compute_loss(ctx, batch, outputs, loss)
return loss
|