mcode-neural 0.5.0__py3-none-any.whl

mcode_neural/train.py

@@ -0,0 +1,756 @@
+# region  F I L E
+# <copyright file="mcode-neural/src/mcode_neural/train.py" company="MicroCODE Incorporated">Copyright © 2026 MicroCODE, Inc. Troy, MI</copyright><author>Timothy J. McGuire</author>
+# region  M O D U L E
+# region  D O C U M E N T A T I O N
+"""
+Project:  MicroCODE Python Packages
+Customer: Internal + MIT xPRO Course
+@module   'mcode_neural.train'
+@memberof mcode_neural
+@created  October 2026
+@author   Timothy McGuire, MicroCODE, Inc.
+@description >
+MicroCODE Shared Training Library
+
+LICENSE:
+--------
+MIT License: MicroCODE.mcode-neural
+
+Copyright (c) 2026 Timothy McGuire, MicroCODE, Inc.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+
+DESCRIPTION:
+------------
+This module implements MicroCODE's training loop for our models, one epoch
+loop for every Model, a single neuron or a PyTorch network, classification,
+regression, and autoencoder: each epoch trains on the model's 'train'
+phase and scores every other phase through the Model's own step, keeping
+the weights of the epoch its kind's rule picks on 'val'.
+
+
+REFERENCES:
+-----------
+1. MIT xPRO Course: AI and Deep Learning: Neural Network Applications
+
+2. PyTorch -- Transfer Learning for Computer Vision Tutorial (the course's 'train_model()')
+   https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html
+
+3. PyTorch -- torch.optim (optimizers and learning rate schedulers)
+   https://pytorch.org/docs/stable/optim.html
+"""
+
+# endregion
+
+# region  I M P O R T S
+
+import inspect
+import numbers
+import time
+
+import copy
+import gc
+import itertools
+import os
+import traceback
+
+import torch
+
+from .data import data_loaders
+
+from .model import HISTORY_DEFAULTS, MODEL_KINDS, _plain
+
+from .model import _check_model
+
+from .network import NETWORK_KINDS
+
+from .model import _out_of_memory
+
+from .model import _gpu_memory, model_save
+
+from .model import _peaks
+
+# endregion
+
+# region  P R I V A T E   F U N C T I O N S
+
+
+def _check_count(name, count):
+    """
+    Raises a ValueError naming 'name' unless 'count' is a whole number of at least 1.
+
+    Args:
+        name:   the argument being checked, e.g. "num_epochs".
+        count:  the value to check; NumPy integers are whole numbers, bools are not.
+    """
+    if isinstance(count, bool) or not isinstance(count, numbers.Integral) or count < 1:
+        raise ValueError(
+            f"'{name}' must be a whole number of at least 1, got {count!r}"
+        )
+
+
+def _check_phases(model):
+    """
+    Raises a ValueError unless the model's data has "train" first, a "val" phase, and labels in every phase.
+
+    Args:
+        model:  the model to train, its data loaded.
+    """
+    datasets = model.data.datasets
+    phases = list(datasets)
+    if not phases or phases[0] != 'train':
+        raise ValueError(
+            f"'model.data' must start with 'train', so each epoch scores the weights it trained; its phases are {phases}"
+        )
+    if 'val' not in phases:
+        raise ValueError(
+            f"'model.data' needs a 'val' phase to pick the best epoch; its phases are {phases}"
+        )
+    # a phase of new rows without labels can be predicted, but not scored every epoch
+    unlabeled = [
+        phase
+        for phase, dataset in datasets.items()
+        if len(dataset) and len(dataset[0]) != 2
+    ]
+    if unlabeled:
+        raise ValueError(
+            f"'model.data' phase(s) {unlabeled} have no labels to train or score against"
+        )
+
+
+def _check_keys(name, value, model):
+    """
+    Raises a TypeError or ValueError naming 'name' unless 'value' is one item, or a dict keyed by 'model.optimizer_keys'.
+
+    Args:
+        name:   the argument being checked, "optimizer" or "scheduler".
+        value:  the caller's optimizer or scheduler.
+        model:  the Model to train; its 'optimizer_keys' are () for one optimizer.
+    """
+    owner = f"'{type(model).__name__}.optimizer_keys'"
+    keys = list(model.optimizer_keys)
+    if not keys:
+        if isinstance(value, dict):
+            raise TypeError(
+                f"'{name}' must be one {name} for a {type(model).__name__}, which names no"
+                f' {owner}, got a dict of {list(value)}'
+            )
+        return
+    if not isinstance(value, dict):
+        raise TypeError(
+            f"'{name}' must be a dict of one {name} per key of {owner} {keys}, got {type(value).__name__}"
+        )
+    missing = [key for key in keys if key not in value]
+    extra = [key for key in value if key not in keys]
+    if missing or extra:
+        raise ValueError(
+            f"'{name}' must have exactly the keys of {owner} {keys}; missing {missing}, extra {extra}"
+        )
+
+
+def _each(value, function):
+    """
+    Returns function(value) for one item, or {key: function(item)} for a dict of them, e.g. one optimizer per network.
+
+    Args:
+        value:      one optimizer or scheduler, or a dict of them keyed by 'optimizer_keys'.
+        function:   applied to each one.
+    """
+    if isinstance(value, dict):
+        return {key: function(item) for key, item in value.items()}
+    return function(value)
+
+
+def _optimizer_record(optimizer):
+    """
+    Returns an optimizer's settings for a "train" event: its name, the constructor's settings given, then its defaults.
+
+    Args:
+        optimizer:  a PyTorch optimizer or a NeuronOptimizer; both hold 'defaults' and 'param_groups'.
+
+    Returns:
+        {"name", the settings given by the constructor's own names, "defaults": {the rest}}, and
+        "param_groups" when a group's setting differs from the constructor's, e.g. a decayed 'lr'.
+    """
+    # the constructor's settings by its own names: given, or left at the default its signature names
+    signature = inspect.signature(type(optimizer)).parameters
+    settings = _plain(optimizer.defaults, 'hyperparams.optimizer')
+    given = {
+        key: value
+        for key, value in settings.items()
+        if key in signature and value != signature[key].default
+    }
+    record = {
+        'name': type(optimizer).__name__,
+        **given,
+        HISTORY_DEFAULTS: {
+            key: value for key, value in settings.items() if key not in given
+        },
+    }
+    # a group's setting only where it differs from the constructor's, e.g. a decayed 'lr'
+    groups = [
+        {
+            key: value
+            for key, value in _plain(
+                {key: group[key] for key in settings},
+                f'hyperparams.optimizer.param_groups[{i}]',
+            ).items()
+            if value != settings[key]
+        }
+        for i, group in enumerate(optimizer.param_groups)
+    ]
+    if any(groups):
+        record['param_groups'] = groups
+    return record
+
+
+def _schedule_record(scheduler):
+    """Returns a scheduler's name and public 'state_dict()' for a "train" event; torch's own '_' keys are its internals."""
+    state = scheduler.state_dict()
+    return {
+        'name': type(scheduler).__name__,
+        **{key: value for key, value in state.items() if not key.startswith('_')},
+    }
+
+
+def _limit_text(device):
+    """
+    Returns the GPU memory this process may use, as a report names it: "the 6.00 GiB cap", or "the GPU's 8.00 GiB".
+
+    Args:
+        device: a CUDA device.
+    """
+    total, cap = _gpu_memory(device)
+    return f'the {cap:.2f} GiB cap' if cap is not None else f"the GPU's {total:.2f} GiB"
+
+
+def _peak_text(peaks, device):
+    """
+    Returns peaks as one line reads them, e.g. "3.48 GiB allocated, 5.07 GiB reserved of the 6.00 GiB cap".
+
+    Args:
+        peaks:  '_peaks()', or results holding its two keys.
+        device: the CUDA device they were measured on.
+    """
+    return (
+        f'{peaks["peak_allocated_gib"]:.2f} GiB allocated,'
+        f' {peaks["peak_reserved_gib"]:.2f} GiB reserved of {_limit_text(device)}'
+    )
+
+
+def _memory(device):
+    """
+    Returns a network run's memory for its "train" event: its device, its cap, and its peaks, all but the device None on a CPU.
+
+    Args:
+        device: the network's device.
+    """
+    if device.type != 'cuda':
+        return {
+            'device': str(device),
+            'cap_gib': None,
+            'peak_allocated_gib': None,
+            'peak_reserved_gib': None,
+        }
+    _, cap = _gpu_memory(device)
+    return {
+        'device': str(device),
+        'cap_gib': None if cap is None else round(cap, 3),
+        **_peaks(device),
+    }
+
+
+def _release(model, device):
+    """
+    Gives back the GPU memory a network no longer needs: every parameter's gradient, and PyTorch's cached blocks.
+
+    Args:
+        model:  the network.
+        device: its device; on a CPU only the gradients are dropped; None, a NumPy neuron's,
+                holds nothing to release.
+    """
+    if device is None:
+        return
+    model.zero_grad(set_to_none=True)
+    if device.type == 'cuda':
+        # a model and its Networks refer to each other, so a discarded one is freed only when collected
+        gc.collect()
+        torch.cuda.empty_cache()
+
+
+def _clear_frames(error):
+    """
+    Clears the local variables of every finished frame in an error's traceback, and in each error it was raised from.
+
+    Args:
+        error:  the error a run raised.
+
+    Note:
+        A Notebook keeps the last traceback, and its frames hold the failed batch's tensors, so the
+        GPU memory they use is freed only when the frames are cleared; '%debug' then shows no locals.
+    """
+    while error is not None:
+        traceback.clear_frames(error.__traceback__)
+        error = error.__context__
+
+
+def _save_whole(model, path):
+    """
+    Saves 'model' with 'model_save()' to 'path' + ".tmp", then replaces 'path' with it, so a crash mid-write keeps the last file.
+
+    Args:
+        model:  the model to save.
+        path:   the checkpoint file.
+    """
+    temporary = f'{os.fspath(path)}.tmp'
+    model_save(model, temporary)
+    os.replace(temporary, path)
+
+
+def _to_cpu(value):
+    """
+    Returns a copy of 'value' with every tensor in it copied to system RAM, e.g. an optimizer's 'state_dict()'.
+
+    Args:
+        value:  a tensor, or a dict, list, or tuple holding tensors and plain values.
+    """
+    if isinstance(value, torch.Tensor):
+        return value.detach().to('cpu', copy=True)
+    if isinstance(value, dict):
+        return {key: _to_cpu(item) for key, item in value.items()}
+    if isinstance(value, (list, tuple)):
+        return type(value)(_to_cpu(item) for item in value)
+    return copy.deepcopy(value)
+
+
+def _run_phase(model, loader, phase, optimizer):
+    """
+    Runs one pass over one phase's loader through the Model's own step, training only in the "train" phase.
+
+    Args:
+        model:      the Model, which moves each batch to its own device.
+        loader:     the phase's DataLoader of (inputs, labels) batches.
+        phase:      the phase's name, its key in 'model.data.datasets'.
+        optimizer:  what 'model.train_step()' steps, once per batch in the "train" phase.
+
+    Returns:
+        {metric: value}, each metric averaged over every sample the loader served.
+    """
+    name = f"'model.data' phase {phase!r}"
+    metrics = MODEL_KINDS[model.kind]['metrics']
+    totals = dict.fromkeys(metrics, 0.0)
+    samples = 0
+
+    for inputs, labels in loader:
+        # S T E P: the Model's own, trained in the "train" phase only, its totals over the batch
+        if phase == 'train':
+            step = model.train_step(inputs, labels, optimizer, name)
+        else:
+            step = model.score_step(inputs, labels, name)
+        # a step that scores other metrics would chart curves its kind does not name
+        if tuple(step) != metrics:
+            raise ValueError(
+                f"'{type(model).__name__}' scored {list(step)} in {name};"
+                f' its kind {model.kind!r} scores {list(metrics)}, in that order'
+            )
+
+        # T A L L Y: each batch's totals, so the phase average weighs every sample alike
+        samples += len(inputs)
+        for metric, total in step.items():
+            totals[metric] += total
+
+    if not samples:
+        raise ValueError(f'{name} served no samples')
+    return {metric: total / samples for metric, total in totals.items()}
+
+
+def _train(model, num_epochs, run_phase, hyperparams, device=None, checkpoint=None):
+    """
+    Runs the epoch loop every Model shares, then keeps the weights of the val epoch its kind's 'is_better' picks last.
+
+    Args:
+        model:          the model, its data loaded and checked by '_check_phases()'.
+        num_epochs:     the number of epochs, already checked.
+        run_phase:      run_phase(phase) -> {metric: value}; trains in the "train" phase only.
+        hyperparams:    the trainer's settings, plain values from '_plain()', for the "train" event.
+        device:         a network's device, whose memory the run records, and on CUDA prints and
+                        releases; None for a single neuron, which records none.
+        checkpoint:     a file the model is saved to at each new best epoch and at the end; None
+                        saves nothing.
+
+    Returns:
+        the same 'model', its best val epoch's weights loaded, with a "train" event appended to
+        its history: 'hyperparams', and the results {"curves", "best_epoch",
+        f"best_val_{metric}", "seconds", "tests": {}}, so 'model.curves' and 'model.best_epoch'
+        read this run; a network's results add "device", "cap_gib", "peak_allocated_gib", and
+        "peak_reserved_gib" after "seconds".
+
+    Note:
+        An interrupted or failed run, e.g. Interrupt Kernel or an out-of-memory error, puts back
+        the weights the model had when it started, records no event, and raises again. A
+        checkpoint written before then keeps its best epoch, its last "train" event the run so
+        far, with "status": "in process".
+    """
+    # S E T   U P: the model kind's metrics, one empty curve per phase and metric, no best epoch yet
+    model_kind = MODEL_KINDS[model.kind]
+    phases = list(model.data.datasets)
+    width = max(len(phase) for phase in phases)
+    training_curves = {
+        f'{phase}_{metric}': [] for phase in phases for metric in model_kind['metrics']
+    }
+    best = None
+    start = time.perf_counter()
+
+    # H I S T O R Y: this run's settings and results, one event, so an earlier run's stay as they were
+    def results(seconds):
+        return {
+            'curves': training_curves,
+            'best_epoch': best_epoch,
+            f'best_val_{model_kind["best"]}': best,
+            'seconds': round(seconds, 3),
+            **({} if device is None else _memory(device)),
+            'tests': {},
+        }
+
+    # S T A R T: the weights before this run, in system RAM, put back if the run is interrupted or fails
+    start_weights = model.weights()
+    try:
+        # E P O C H   L O O P: train, then score every phase, keeping the best val epoch's weights
+        for epoch in range(num_epochs):
+            print(f'\nEpoch {epoch}/{num_epochs - 1}')
+            print('-' * 10)
+
+            for phase in phases:
+                values = run_phase(phase)
+                for metric, value in values.items():
+                    training_curves[f'{phase}_{metric}'].append(value)
+                scores = ' '.join(
+                    f'{metric.title()}: {value:.4f}' for metric, value in values.items()
+                )
+                print(f'{phase:{width}} {scores}')
+
+                if phase == 'val':
+                    score = values[model_kind['best']]
+                    if best is None or model_kind['is_better'](score, best):
+                        best = score
+                        best_epoch = epoch
+                        best_weights = model.weights()
+
+            # C H E C K P O I N T: at an epoch's end, a new best is saved with the run so far, in the file only
+            if checkpoint is not None and best_epoch == epoch:
+                so_far = results(time.perf_counter() - start)
+                model._record('train', hyperparams, so_far | {'status': 'in process'})
+                try:
+                    _save_whole(model, checkpoint)
+                finally:
+                    model.history.pop()
+    except BaseException:
+        model.load_weights(start_weights)
+        raise
+
+    # B E S T   E P O C H: report the run's time and best score, then restore that epoch's weights
+    elapsed = time.perf_counter() - start
+    minutes, seconds = divmod(round(elapsed), 60)
+    model.load_weights(best_weights)
+    finished = results(elapsed)
+    print(f'\nTraining complete in {minutes}m {seconds}s')
+    if device is not None and device.type == 'cuda':
+        print(f'Peak GPU memory: {_peak_text(finished, device)}')
+        _release(model, device)
+        held = torch.cuda.memory_reserved(device) / 2**30
+        print(f'GPU memory held: {held:.2f} GiB')
+    print(
+        f'\nBEST val {model_kind["best"].title()}: {best:.4f} at epoch {best_epoch}\n\n'
+    )
+
+    model._record('train', hyperparams, finished)
+    if checkpoint is not None:
+        _save_whole(model, checkpoint)
+
+    return model
+
+
+# endregion
+
+# region  P U B L I C   F U N C T I O N S
+
+
+def train_model(model, optimizer, scheduler, num_epochs, batch_size, checkpoint=None):
+    """
+    Trains a Model on its own data, scoring every phase each epoch, and keeps its best val epoch's weights.
+
+    Args:
+        model:      the Model to train, its data loaded, e.g. ClassificationModel,
+                    BinaryClassificationModel, RegressionModel, AutoencoderModel, or a single
+                    neuron such as LogisticClassifierNeuronModel; its own 'train_step()' and
+                    'score_step()' make its losses, step its optimizer, and score each batch.
+        optimizer:  the optimizer of 'model.parameters()', e.g. torch.optim.Adam; a
+                    NeuronOptimizer for a single neuron; or, for a Model that names
+                    'optimizer_keys', a dict of one optimizer per key, e.g.
+                    {"generator": ..., "discriminator": ...}.
+        scheduler:  steps once per epoch, after the "train" phase, e.g. ExponentialLR; None for
+                    no schedule. For a Model that names 'optimizer_keys', a dict of one
+                    scheduler per key. PyTorch's schedulers take only PyTorch optimizers.
+        num_epochs: the number of epochs, a whole number of at least 1.
+        batch_size: the number of samples in each batch, a whole number of at least 1.
+        checkpoint: a file path, e.g. "../Models/vgg11.pt", that the model is saved to with
+                    'model_save()' at each epoch whose val score is the best so far, and at the
+                    end; its folder must exist. None, the default, saves nothing.
+
+    Returns:
+        the same 'model', its best val epoch's weights loaded, a network in eval() mode, with
+        'model.curves' = {f"{phase}_{metric}": [one float per epoch]} and that epoch's number
+        in 'model.best_epoch', as 'show_training_curves(model)' reads them, both from the
+        "train" event this run appends to 'model.history'.
+
+    Note:
+        The phases are those of 'model.data': "train" first, then "val", which picks the best
+        epoch, and any others, e.g. "test", scored every epoch; each is served by an unshuffled
+        DataLoader of 'batch_size', as 'data_loaders()' makes.
+        A classifier trains on nn.CrossEntropyLoss of its logits, scores "loss" and "acc", and
+        keeps its highest val acc; a regressor trains on nn.MSELoss, scores "loss" only, its
+        (N, 1) output flattened to (N,) targets, and keeps its lowest val loss.
+        A binary classifier trains on nn.BCEWithLogitsLoss: its labels, (N,) for one output or
+        (N, k) for k, must be 0 or 1 and are cast to the model's float dtype, so integer labels
+        train as they are; it predicts 'logit > 0', its "acc" is the share of (sample, label)
+        answers that are right, and it keeps its highest val acc.
+        An autoencoder trains on nn.MSELoss of each sample against itself, scores
+        "loss" only, and keeps its lowest val loss; its labels only serve the batches, and every
+        input must lie in its output style's 'output_range'.
+        A single neuron trains with a NeuronOptimizer on its own loss ('loss()' and
+        'loss_gradient()'): each "train" batch is scored, then the neuron steps by the mean of
+        its samples' updates, so 'batch_size=1' is the course's per-sample loop. Its labels
+        must be one per sample, each one of its 'label_value'. A classifier neuron scores
+        "loss" and "acc" and keeps its highest val acc, predicting its higher label where a is
+        above its threshold; a regressor neuron scores "loss" only and keeps its lowest val
+        loss. A neuron has no device, so its results record no memory.
+        A Model of several networks, e.g. a GAN, names its optimizers' keys in
+        'optimizer_keys', and a missing or extra key fails naming it; its step returns one
+        total per metric of its kind, e.g. "dloss" and "gloss", and a kind whose 'is_better'
+        always holds keeps its last epoch.
+        The first epoch always counts as the best so far; a later epoch replaces it only by
+        scoring strictly better, so the earliest of equal epochs is kept. Each batch is moved
+        to the model's own device and shaped as the network reads each sample, its 'in_shape':
+        one row for a Linear network, one (channels, height, width) image for a network with
+        conv stages. Epochs are numbered from 0 in every printed line.
+        The "train" event records the optimizer's name and its constructor's settings by their
+        own names, e.g. "lr": the ones given, then those left at their defaults under
+        "defaults", which 'show_model_history()' marks "(default)"; a setting the optimizer
+        class fixes itself, e.g. AdamW's 'decoupled_weight_decay', reads as a default. A param
+        group's setting is recorded, under "param_groups", only where it differs from the
+        constructor's when the run starts, e.g. a decayed 'lr'. It also records the scheduler's
+        name and public 'state_dict()' (or None), 'num_epochs', and 'batch_size'; training
+        again appends another event, so every run is kept. A dict of optimizers is recorded
+        by its keys, {key: record}, and a dict of schedulers the same way.
+        A batch that runs past the GPU memory this process may use, the cap 'model_device()'
+        set or the GPU's whole memory, raises torch.OutOfMemoryError naming the phase and
+        'batch_size', e.g. "phase 'train' at batch_size=100 needs more than the 6.00 GiB cap;
+        try a smaller batch_size"; PyTorch's own error is its cause.
+        The "train" event's results also record the run's "device", e.g. "cuda:0", the cap
+        'model_device()' set in GiB ("cap_gib", None with no cap), and its peak GPU memory in
+        GiB, "peak_allocated_gib" and "peak_reserved_gib" (None on a CPU). On CUDA it prints,
+        after the run's time, "Peak GPU memory: 3.48 GiB allocated, 5.07 GiB reserved of the
+        6.00 GiB cap" and, once released, "GPU memory held: 0.55 GiB".
+        However the run ends, every parameter's gradient is dropped and PyTorch's cached GPU
+        blocks are returned to the driver. An interrupted or failed run, e.g. Interrupt Kernel
+        or an out-of-memory error, records no event and puts back the weights the model had when
+        it started; its traceback's local variables are cleared, so the failed batch leaves the
+        GPU ('%debug' shows no locals). The optimizer's state, e.g. Adam's two moments per
+        weight, is the caller's: free it with 'del optimizer', or build a new optimizer for the
+        next phase.
+        A checkpoint is written to "<checkpoint>.tmp" and then replaces the file, so a crash
+        mid-write keeps the last best. Until the run ends its history's last "train" event is
+        the run so far, with "status": "in process", which 'show_model_history()' reads "In
+        process..."; the finished run rewrites it with the finished event. 'model_load()' reads
+        it as any saved model.
+
+    Example:
+        >>> model.data = data_frame_split(df, input_columns, "color", torch.int64,
+        ...                               input_columns, wine_classes, seed=42)
+        >>> optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
+        >>> scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.95)
+        >>> model = train_model(model, optimizer, scheduler, num_epochs=50, batch_size=10)
+        >>> show_training_curves(model)
+        >>> neuron = train_model(neuron, NeuronOptimizer(lr=0.01), None, num_epochs=10, batch_size=1)
+    """
+    # V A L I D A T E: the model, its phases, the counts, and its optimizers fail by name before training
+    _check_model(model)
+    _check_phases(model)
+    _check_count('num_epochs', num_epochs)
+    _check_count('batch_size', batch_size)
+    _check_keys('optimizer', optimizer, model)
+    if scheduler is not None:
+        _check_keys('scheduler', scheduler, model)
+    if checkpoint is not None:
+        if not isinstance(checkpoint, (str, os.PathLike)):
+            raise TypeError(
+                f'\'checkpoint\' must be a file path, e.g. "../Models/vgg11.pt", got {type(checkpoint).__name__}'
+            )
+        folder = os.path.dirname(os.path.abspath(checkpoint))
+        if not os.path.isdir(folder):
+            raise ValueError(
+                f"'checkpoint' {os.fspath(checkpoint)!r} is in a folder that does not exist: {folder!r}"
+            )
+
+    # H Y P E R P A R A M E T E R S: this run's settings, made plain now, so one a saved model
+    # cannot keep fails by name before training; torch's own '_' keys are its internals
+    hyperparams = _plain(
+        {
+            'trainer': 'train_model',
+            'optimizer': _each(optimizer, _optimizer_record),
+            'scheduler': None
+            if scheduler is None
+            else _each(scheduler, _schedule_record),
+            'num_epochs': num_epochs,
+            'batch_size': batch_size,
+            'checkpoint': None if checkpoint is None else os.fspath(checkpoint),
+        },
+        'hyperparams',
+    )
+
+    # S E T   U P: one loader per phase, and the device its network kind trains on, None for NumPy
+    loaders = data_loaders(model.data.datasets, batch_size)
+    device = NETWORK_KINDS[model.network_kind]['device'](model)
+
+    def run_phase(phase):
+        try:
+            values = _run_phase(model, loaders[phase], phase, optimizer)
+        except torch.OutOfMemoryError as error:
+            raise _out_of_memory(phase, batch_size, device) from error
+        if phase == 'train' and scheduler is not None:
+            _each(scheduler, lambda item: item.step())
+        return values
+
+    # M E M O R Y: the peak is this run's; whatever ends it, its gradients and cached blocks go back
+    if device is not None and device.type == 'cuda':
+        torch.cuda.reset_peak_memory_stats(device)
+    try:
+        return _train(model, num_epochs, run_phase, hyperparams, device, checkpoint)
+    except BaseException as error:
+        _clear_frames(error)
+        raise
+    finally:
+        _release(model, device)
+
+
+def train_probe(model, optimizer, batch_sizes, num_batches=2):
+    """
+    Measures the peak GPU memory of training a network at each batch size, smallest first, and leaves it as it was.
+
+    Args:
+        model:          the network to probe, on a CUDA device, its data loaded, e.g. ClassificationModel.
+        optimizer:      the optimizer the run will use, of 'model.parameters()', or the dict of
+                        them a Model that names 'optimizer_keys' takes; it steps, so its state,
+                        e.g. Adam's two moments per weight, counts in each peak.
+        batch_sizes:    the batch sizes to try, e.g. [32, 64, 100], each a whole number of at least 1
+                        and at most the size of the "train" phase.
+        num_batches:    the training batches run at each size, a whole number of at least 1.
+
+    Returns:
+        {batch_size: {"peak_allocated_gib", "peak_reserved_gib"}}, smallest first, in GiB to 3
+        places; the first size that does not fit maps to None, and no larger size is tried.
+
+    Note:
+        Each size trains 'num_batches' batches of the "train" phase as 'train_model()' does,
+        through the Model's own 'train_step()'. Reserved is what the cap of
+        'model_device()' bounds. One line is printed per size, e.g. "batch_size  64: 5.00 GiB
+        allocated, 5.56 GiB reserved of the 6.00 GiB cap" or "batch_size 100: does not fit the
+        6.00 GiB cap". Afterwards the model's weights, the optimizer's state, and torch's random
+        state are put back, its gradients and PyTorch's cached GPU blocks are released, and
+        nothing is recorded in its history; the scheduler is not used.
+
+    Example:
+        >>> model = ClassificationModel(...).to(model_device())
+        >>> optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
+        >>> train_probe(model, optimizer, [32, 48, 64, 100])
+        >>> model = train_model(model, optimizer, None, num_epochs=10, batch_size=64)
+    """
+    # V A L I D A T E: the network, its data, the sizes, and its device fail by name before any batch
+    _check_model(model)
+    _check_phases(model)
+    _check_keys('optimizer', optimizer, model)
+    if not isinstance(batch_sizes, (list, tuple)) or not batch_sizes:
+        raise ValueError(
+            f"'batch_sizes' must be a list of whole numbers, e.g. [32, 64, 100], got {batch_sizes!r}"
+        )
+    dataset = model.data.datasets['train']
+    for i, batch_size in enumerate(batch_sizes):
+        _check_count(f'batch_sizes[{i}]', batch_size)
+        if batch_size > len(dataset):
+            raise ValueError(
+                f"'batch_sizes[{i}]' is {batch_size}, more than the {len(dataset)} samples of the 'train' phase"
+            )
+    _check_count('num_batches', num_batches)
+    device = NETWORK_KINDS[model.network_kind]['device'](model)
+    if device is None:
+        raise ValueError(
+            f"'model' is a {type(model).__name__}, which has no device; 'train_probe()' measures"
+            ' the GPU memory of a PyTorch network'
+        )
+    if device.type != 'cuda':
+        raise ValueError(
+            f"'model' is on the {device}; 'train_probe()' measures GPU memory, so move it to the GPU"
+            " first, e.g. 'model.to(model_device())'"
+        )
+
+    # S T A R T: what the probe puts back, kept in system RAM so it adds nothing to a peak
+    start_weights = model.weights()
+    # each optimizer with the state it is put back to
+    start_states = _each(optimizer, lambda item: (item, _to_cpu(item.state_dict())))
+    was_training = model.training
+    width = len(str(max(batch_sizes)))
+    peaks = {}
+
+    # P R O B E: each size from a released GPU, its peak alone; the first that does not fit ends it
+    try:
+        with torch.random.fork_rng(devices=[device.index]):
+            for batch_size in sorted(batch_sizes):
+                _release(model, device)
+                torch.cuda.reset_peak_memory_stats(device)
+                loader = data_loaders({'train': dataset}, batch_size)['train']
+                try:
+                    _run_phase(
+                        model,
+                        itertools.islice(loader, num_batches),
+                        'train',
+                        optimizer,
+                    )
+                except torch.OutOfMemoryError as error:
+                    _clear_frames(error)
+                    peaks[batch_size] = None
+                    print(
+                        f'batch_size {batch_size:>{width}}: does not fit {_limit_text(device)}'
+                    )
+                    break
+                peaks[batch_size] = _peaks(device)
+                print(
+                    f'batch_size {batch_size:>{width}}: {_peak_text(peaks[batch_size], device)}'
+                )
+    finally:
+        model.load_weights(start_weights)
+        _each(start_states, lambda start: start[0].load_state_dict(start[1]))
+        model.train(was_training)
+        _release(model, device)
+    return peaks
+
+
+# endregion
+# endregion
+# endregion