Skip to content

traingen

package license python runtime docs

traingen contains shared training utilities for design-generators packages that train models locally. It keeps reusable PyTorch Lightning CLI integration points out of individual model packages while leaving model-specific data modules, losses, and training loops in models/*.

Use traingen for generic training infrastructure such as LightningCLI helpers. Use traingen-parity for deterministic trace capture and reference/target comparison reports.

traingen.optim.KerasAdam implements the update order and epsilon placement used by Keras 2 Adam. CanvasVAE and PixelVAE use it when comparing TensorFlow training updates with their PyTorch models.

Install

uv sync --package traingen

Install from outside the workspace with pip's direct-reference subdirectory form:

pip install "traingen @ git+https://github.com/creative-graphic-design/design-generators.git#subdirectory=lib/traingen"

API Entry Points

Run a model-agnostic LightningCLI from a training YAML that declares model.class_path and data.class_path.

uv run --package layout-flow --extra training \
  traingen fit \
  --config models/layout-flow/configs/training/smoke.yaml

Create package-local LightningCLI subclasses only when a model package needs custom CLI arguments or defaults that cannot live in YAML.

uv run --package traingen python

from traingen.lightning.cli import lightning_cli_class


class MyLightningCLI(lightning_cli_class()):
    pass

Use traingen.lightning.steps for generic Lightning loss reduction and logging in model training steps.

from traingen.lightning.steps import finish_training_step, log_validation_loss


def training_step(self, batch, batch_idx):
    losses = self.compute_losses(batch)
    total, self.latest_step_trace = finish_training_step(self, losses, {})
    return total


def validation_step(self, batch, batch_idx):
    total = self.compute_validation_loss(batch)
    log_validation_loss(self, total)
    return total

Scope

  • Keep reusable training utilities here only when more than one package that trains models locally can use them.
  • Keep dataset-specific modules, model losses, condition policies, and checkpoint conversion in each model package.
  • Keep deterministic comparison logic in traingen-parity so regular training helpers do not depend on parity-only behavior.

Pointers

API Reference

Shared training utilities for design generator packages.

lightning

Optional Lightning helpers for train-ourselves packages.

cli

Shared LightningCLI entry point for train-ourselves packages.

lightning_cli_class 

lightning_cli_class() -> type[LightningCLI]

Return LightningCLI without importing Lightning at package import time.

Source code in lib/traingen/src/traingen/lightning/cli.py 
11
12
13
14
15
def lightning_cli_class() -> type[LightningCLI]:
    """Return LightningCLI without importing Lightning at package import time."""
    from lightning.pytorch.cli import LightningCLI

    return LightningCLI

main 

main(args: list[str] | None = None) -> None

Run a model-agnostic LightningCLI from YAML class_path entries.

Parameters:

Name Type Description Default
args list[str] | None

Optional CLI arguments for tests. When omitted, LightningCLI reads sys.argv.

 None

Raises:

Type Description
SystemExit

If LightningCLI argument parsing fails.
Source code in lib/traingen/src/traingen/lightning/cli.py 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
def main(args: list[str] | None = None) -> None:
    """Run a model-agnostic LightningCLI from YAML ``class_path`` entries.

    Args:
        args: Optional CLI arguments for tests. When omitted, LightningCLI reads
            ``sys.argv``.

    Raises:
        SystemExit: If LightningCLI argument parsing fails.
    """
    lightning_cli_class()(
        model_class=None,
        datamodule_class=None,
        subclass_mode_model=True,
        subclass_mode_data=True,
        args=args,
    )

steps

Generic Lightning training-step logging and reduction helpers.

ScalarLogger

Bases: Protocol

Minimal scalar logging protocol implemented by Lightning modules.

Source code in lib/traingen/src/traingen/lightning/steps.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
class ScalarLogger(Protocol):
    """Minimal scalar logging protocol implemented by Lightning modules."""

    def log(
        self,
        name: str,
        value: Float[torch.Tensor, ""],
        *,
        prog_bar: bool = False,
        on_step: bool | None = None,
        on_epoch: bool | None = None,
        batch_size: int | None = None,
    ) -> None:
        """Log a scalar training metric.

        Args:
            name: Metric name.
            value: Scalar tensor value.
            prog_bar: Whether to show the value in the progress bar.
            on_step: Whether to aggregate the value per step.
            on_epoch: Whether to aggregate the value per epoch.
            batch_size: Batch size used by Lightning metric aggregation.
        """
        _ = (name, value, prog_bar, on_step, on_epoch, batch_size)
log 
log(
    name: str,
    value: Float[Tensor, ""],
    *,
    prog_bar: bool = False,
    on_step: bool | None = None,
    on_epoch: bool | None = None,
    batch_size: int | None = None,
) -> None

Log a scalar training metric.

Parameters:

Name Type Description Default
name str

Metric name.

 required
value Float[Tensor, '']

Scalar tensor value.

 required
prog_bar bool

Whether to show the value in the progress bar.

 False
on_step bool | None

Whether to aggregate the value per step.

 None
on_epoch bool | None

Whether to aggregate the value per epoch.

 None
batch_size int | None

Batch size used by Lightning metric aggregation.

 None
Source code in lib/traingen/src/traingen/lightning/steps.py 
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
def log(
    self,
    name: str,
    value: Float[torch.Tensor, ""],
    *,
    prog_bar: bool = False,
    on_step: bool | None = None,
    on_epoch: bool | None = None,
    batch_size: int | None = None,
) -> None:
    """Log a scalar training metric.

    Args:
        name: Metric name.
        value: Scalar tensor value.
        prog_bar: Whether to show the value in the progress bar.
        on_step: Whether to aggregate the value per step.
        on_epoch: Whether to aggregate the value per epoch.
        batch_size: Batch size used by Lightning metric aggregation.
    """
    _ = (name, value, prog_bar, on_step, on_epoch, batch_size)

sum_loss_values 

sum_loss_values(
    losses: Mapping[str, Float[Tensor, ""]],
) -> Float[torch.Tensor, ""]

Sum scalar loss values with the canonical training reduction.

Parameters:

Name Type Description Default
losses Mapping[str, Float[Tensor, '']]

Mapping from metric names to scalar loss tensors.

 required

Returns:

Type Description
Float[Tensor, '']

Scalar tensor containing the sum of all loss values.

Examples:

>>> import torch
>>> sum_loss_values({"a": torch.tensor(1.0), "b": torch.tensor(2.0)})
tensor(3.)
Source code in lib/traingen/src/traingen/lightning/steps.py 
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
def sum_loss_values(
    losses: Mapping[str, Float[torch.Tensor, ""]],
) -> Float[torch.Tensor, ""]:
    """Sum scalar loss values with the canonical training reduction.

    Args:
        losses: Mapping from metric names to scalar loss tensors.

    Returns:
        Scalar tensor containing the sum of all loss values.

    Examples:
        >>> import torch
        >>> sum_loss_values({"a": torch.tensor(1.0), "b": torch.tensor(2.0)})
        tensor(3.)
    """
    import torch

    return torch.stack(tuple(losses.values())).sum()

log_training_losses 

log_training_losses(
    logger: ScalarLogger,
    losses: Mapping[str, Float[Tensor, ""]],
    total: Float[Tensor, ""],
    *,
    batch_size: int = 1,
) -> None

Log per-component and total training losses.

Parameters:

Name Type Description Default
logger ScalarLogger

Object exposing Lightning-compatible log.

 required
losses Mapping[str, Float[Tensor, '']]

Per-component scalar loss values.

 required
total Float[Tensor, '']

Total scalar training loss.

 required
batch_size int

Batch size used by Lightning metric aggregation.

 1

Returns:

Type Description
None

None.
Source code in lib/traingen/src/traingen/lightning/steps.py 
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
def log_training_losses(
    logger: ScalarLogger,
    losses: Mapping[str, Float[torch.Tensor, ""]],
    total: Float[torch.Tensor, ""],
    *,
    batch_size: int = 1,
) -> None:
    """Log per-component and total training losses.

    Args:
        logger: Object exposing Lightning-compatible ``log``.
        losses: Per-component scalar loss values.
        total: Total scalar training loss.
        batch_size: Batch size used by Lightning metric aggregation.

    Returns:
        None.
    """
    for key, value in losses.items():
        logger.log(key, value, on_step=True, on_epoch=True, batch_size=batch_size)

    logger.log(
        "train_loss",
        total,
        prog_bar=True,
        on_step=True,
        on_epoch=True,
        batch_size=batch_size,
    )

finish_training_step 

finish_training_step(
    logger: ScalarLogger,
    losses: Mapping[str, Float[Tensor, ""]],
    trace: Mapping[str, Shaped[Tensor, "..."]],
    *,
    batch_size: int = 1,
) -> tuple[
    Float[torch.Tensor, ""],
    dict[str, Shaped[torch.Tensor, "..."]],
]

Reduce losses, log training metrics, and append train_loss to a trace.

Parameters:

Name Type Description Default
logger ScalarLogger

Object exposing Lightning-compatible log.

 required
losses Mapping[str, Float[Tensor, '']]

Per-component scalar loss values.

 required
trace Mapping[str, Shaped[Tensor, '...']]

Training trace entries produced before the optimizer step.

 required
batch_size int

Batch size used by Lightning metric aggregation.

 1

Returns:

Type Description
tuple[Float[Tensor, ''], dict[str, Shaped[Tensor, '...']]]

Total scalar loss and an updated detached trace mapping.
Source code in lib/traingen/src/traingen/lightning/steps.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
def finish_training_step(
    logger: ScalarLogger,
    losses: Mapping[str, Float[torch.Tensor, ""]],
    trace: Mapping[str, Shaped[torch.Tensor, "..."]],
    *,
    batch_size: int = 1,
) -> tuple[Float[torch.Tensor, ""], dict[str, Shaped[torch.Tensor, "..."]]]:
    """Reduce losses, log training metrics, and append ``train_loss`` to a trace.

    Args:
        logger: Object exposing Lightning-compatible ``log``.
        losses: Per-component scalar loss values.
        trace: Training trace entries produced before the optimizer step.
        batch_size: Batch size used by Lightning metric aggregation.

    Returns:
        Total scalar loss and an updated detached trace mapping.
    """
    total = sum_loss_values(losses)
    log_training_losses(logger, losses, total, batch_size=batch_size)
    return total, {**trace, "train_loss": total.detach()}

log_validation_loss 

log_validation_loss(
    logger: ScalarLogger,
    total: Float[Tensor, ""],
    *,
    batch_size: int = 1,
) -> None

Log the canonical validation loss metric.

Parameters:

Name Type Description Default
logger ScalarLogger

Object exposing Lightning-compatible log.

 required
total Float[Tensor, '']

Scalar validation loss.

 required
batch_size int

Batch size used by Lightning metric aggregation.

 1

Returns:

Type Description
None

None.
Source code in lib/traingen/src/traingen/lightning/steps.py 
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
def log_validation_loss(
    logger: ScalarLogger,
    total: Float[torch.Tensor, ""],
    *,
    batch_size: int = 1,
) -> None:
    """Log the canonical validation loss metric.

    Args:
        logger: Object exposing Lightning-compatible ``log``.
        total: Scalar validation loss.
        batch_size: Batch size used by Lightning metric aggregation.

    Returns:
        None.
    """
    logger.log(
        "val_loss",
        total,
        prog_bar=True,
        on_step=False,
        on_epoch=True,
        batch_size=batch_size,
    )

optim

Optimizers with update semantics shared by training reproduction runs.

KerasAdam

Bases: Optimizer

Adam with the update order and epsilon placement of Keras 2 Adam.

For step t the update is m += (g - m) * (1 - beta1), v += (g**2 - v) * (1 - beta2), and p -= m * alpha / (sqrt(v) + epsilon) with alpha = lr * sqrt(1 - beta2**t) / (1 - beta1**t); all scalars are float32. torch.optim.Adam instead adds epsilon after bias-correcting sqrt(v).

Parameters:

Name Type Description Default
params Iterable[Parameter]

Parameters to optimize.

 required
lr float

Learning rate.

 0.001
betas tuple[float, float]

First and second moment decay rates.

 (0.9, 0.999)
eps float

Denominator epsilon.

 1e-07

Examples:

>>> weight = nn.Parameter(torch.ones(2))
>>> optimizer = KerasAdam([weight], lr=0.1)
>>> weight.grad = torch.ones(2)
>>> optimizer.step()
>>> weight.detach().tolist()
[0.8999999761581421, 0.8999999761581421]
Source code in lib/traingen/src/traingen/optim.py 
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
class KerasAdam(torch.optim.Optimizer):
    """Adam with the update order and epsilon placement of Keras 2 ``Adam``.

    For step ``t`` the update is ``m += (g - m) * (1 - beta1)``,
    ``v += (g**2 - v) * (1 - beta2)``, and
    ``p -= m * alpha / (sqrt(v) + epsilon)`` with
    ``alpha = lr * sqrt(1 - beta2**t) / (1 - beta1**t)``; all scalars are
    float32. ``torch.optim.Adam`` instead adds epsilon after bias-correcting
    ``sqrt(v)``.

    Args:
        params: Parameters to optimize.
        lr: Learning rate.
        betas: First and second moment decay rates.
        eps: Denominator epsilon.

    Examples:
        >>> weight = nn.Parameter(torch.ones(2))
        >>> optimizer = KerasAdam([weight], lr=0.1)
        >>> weight.grad = torch.ones(2)
        >>> optimizer.step()
        >>> weight.detach().tolist()
        [0.8999999761581421, 0.8999999761581421]
    """

    def __init__(
        self,
        params: Iterable[nn.Parameter],
        lr: float = 1e-3,
        betas: tuple[float, float] = (0.9, 0.999),
        eps: float = 1e-7,
    ) -> None:
        """Register parameter groups and hyperparameters."""
        super().__init__(params, {"lr": lr, "betas": betas, "eps": eps})

    @torch.no_grad()
    def step(  # ty: ignore[invalid-method-override]
        self, closure: Callable[[], float | Float[torch.Tensor, ""]] | None = None
    ) -> float | Float[torch.Tensor, ""] | None:
        """Apply one update to every parameter with a gradient."""
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        for group in self.param_groups:
            beta1, beta2 = group["betas"]
            for param in group["params"]:
                if param.grad is None:
                    continue

                state = self.state[param]
                if not state:
                    state["step"] = 0
                    state["exp_avg"] = torch.zeros_like(param)
                    state["exp_avg_sq"] = torch.zeros_like(param)

                state["step"] += 1
                step = torch.tensor(
                    float(state["step"]), dtype=param.dtype, device=param.device
                )
                beta1_power = torch.pow(
                    torch.tensor(beta1, dtype=param.dtype, device=param.device), step
                )
                beta2_power = torch.pow(
                    torch.tensor(beta2, dtype=param.dtype, device=param.device), step
                )
                lr = torch.tensor(group["lr"], dtype=param.dtype, device=param.device)
                alpha = lr * torch.sqrt(1 - beta2_power) / (1 - beta1_power)

                grad = param.grad
                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
                exp_avg.add_((grad - exp_avg) * (1 - beta1))
                exp_avg_sq.add_((grad.square() - exp_avg_sq) * (1 - beta2))
                param.sub_((exp_avg * alpha) / (exp_avg_sq.sqrt() + group["eps"]))

        return loss

__init__ 

__init__(
    params: Iterable[Parameter],
    lr: float = 0.001,
    betas: tuple[float, float] = (0.9, 0.999),
    eps: float = 1e-07,
) -> None

Register parameter groups and hyperparameters.

Source code in lib/traingen/src/traingen/optim.py 
36
37
38
39
40
41
42
43
44
def __init__(
    self,
    params: Iterable[nn.Parameter],
    lr: float = 1e-3,
    betas: tuple[float, float] = (0.9, 0.999),
    eps: float = 1e-7,
) -> None:
    """Register parameter groups and hyperparameters."""
    super().__init__(params, {"lr": lr, "betas": betas, "eps": eps})

step 

step(
    closure: Callable[[], float | Float[Tensor, ""]]
    | None = None,
) -> float | Float[torch.Tensor, ""] | None

Apply one update to every parameter with a gradient.

Source code in lib/traingen/src/traingen/optim.py 
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
@torch.no_grad()
def step(  # ty: ignore[invalid-method-override]
    self, closure: Callable[[], float | Float[torch.Tensor, ""]] | None = None
) -> float | Float[torch.Tensor, ""] | None:
    """Apply one update to every parameter with a gradient."""
    loss = None
    if closure is not None:
        with torch.enable_grad():
            loss = closure()

    for group in self.param_groups:
        beta1, beta2 = group["betas"]
        for param in group["params"]:
            if param.grad is None:
                continue

            state = self.state[param]
            if not state:
                state["step"] = 0
                state["exp_avg"] = torch.zeros_like(param)
                state["exp_avg_sq"] = torch.zeros_like(param)

            state["step"] += 1
            step = torch.tensor(
                float(state["step"]), dtype=param.dtype, device=param.device
            )
            beta1_power = torch.pow(
                torch.tensor(beta1, dtype=param.dtype, device=param.device), step
            )
            beta2_power = torch.pow(
                torch.tensor(beta2, dtype=param.dtype, device=param.device), step
            )
            lr = torch.tensor(group["lr"], dtype=param.dtype, device=param.device)
            alpha = lr * torch.sqrt(1 - beta2_power) / (1 - beta1_power)

            grad = param.grad
            exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
            exp_avg.add_((grad - exp_avg) * (1 - beta1))
            exp_avg_sq.add_((grad.square() - exp_avg_sq) * (1 - beta2))
            param.sub_((exp_avg * alpha) / (exp_avg_sq.sqrt() + group["eps"]))

    return loss