diff --git a/pytensor_ml/optim/train.py b/pytensor_ml/optim/train.py index 386a8fc..d3add27 100644 --- a/pytensor_ml/optim/train.py +++ b/pytensor_ml/optim/train.py @@ -1,10 +1,10 @@ from collections.abc import Sequence -from pytensor.compile import Function, SharedVariable +from pytensor.compile import Function from pytensor.graph.basic import Variable, equal_computations from pytensor.tensor import TensorVariable -from pytensor_ml.optim.base import Parameter, UpdateRule, require_unique_state_names +from pytensor_ml.optim.base import Parameter, UpdateRule, Updates, require_unique_state_names from pytensor_ml.pytensorf import ( collect_clock_updates, collect_data_inputs, @@ -21,14 +21,16 @@ def compile_train( parameters: Sequence[Parameter] | None = None, inputs: Sequence[Variable] | None = None, extra_outputs: Sequence[Variable] | None = None, + extra_updates: Updates | None = None, compile_kwargs: dict | None = None, ) -> Function: """ Compile a one-step training function from a loss graph and an update rule. Differentiates the loss via ``rule``, applies the resulting updates, folds in any non-trainable state - updates (such as batch-norm running statistics), advances every training clock the step reads, and - compiles. The parameters and data inputs are collected from ``loss`` unless given explicitly. + updates (such as batch-norm running statistics) and any given in ``extra_updates``, advances every + training clock the step reads, and compiles. The parameters and data inputs are collected from ``loss`` + unless given explicitly. Parameters ---------- @@ -42,17 +44,25 @@ def compile_train( is still initialized and checkpointed, since :func:`collect_trainable_params` reaches it; only the optimizer skips it. inputs : sequence of Variable, optional - Data inputs of the compiled function, in call order. Collected from ``loss`` and ``extra_outputs`` - with :func:`collect_data_inputs` when omitted; pass them explicitly when call order matters (e.g. - features before targets). + Data inputs of the compiled function, in call order. Collected from ``loss``, ``extra_outputs`` and + ``extra_updates`` with :func:`collect_data_inputs` when omitted; pass them explicitly when call order + matters (e.g. features before targets). extra_outputs : sequence of Variable, optional Diagnostics to return alongside the loss, such as gradient norms or a batch accuracy. Evaluated in the same pass as the gradients, so they see the pre-update parameter values, and they add no non-trainable state updates of their own. A random node reached only through an extra output does still advance its generator, since :func:`~pytensor_ml.pytensorf.function` threads the next-RNG update for every generator the outputs draw from. + extra_updates : dict, optional + State the step should write that no gradient produces -- a target-network sync, a Polyak average, + replay priorities. Mapping from shared variable to its next value, folded in alongside the rule's own + updates. Raise if the rule or the model already writes one of these variables, since two writes to one + variable cannot both take effect. An expression here is part of the step like any other, so a clock it + reads still advances once, and a generator it draws from still advances. compile_kwargs : dict, optional - Extra keyword arguments forwarded to the function compiler. + Extra keyword arguments forwarded to the function compiler. An ``updates`` entry is taken as + ``extra_updates`` rather than forwarded, since the compiler's own ``updates`` carries the whole + assembled step. Raise if a variable is given in both. Returns ------- @@ -61,19 +71,45 @@ def compile_train( ``(loss, *extra_outputs)`` when diagnostics were requested. """ extra_outputs = list(extra_outputs or []) + extra_updates = dict(extra_updates or {}) + # Copied rather than mutated, so popping the caller's `updates` out does not empty their dict. + compile_kwargs = dict(compile_kwargs or {}) + + # `updates` is pytensor's own name for this, so a caller who puts one in compile_kwargs is asking for + # what extra_updates does; take it as one instead of letting it collide with this function's own call. + for variable, new_value in compile_kwargs.pop("updates", {}).items(): + if variable in extra_updates: + raise ValueError( + f"The update for {variable.name!r} is given twice, once in `extra_updates` and once in " + "`compile_kwargs['updates']`. Both are folded into the training step, so keep whichever one " + "is right and drop the other." + ) + extra_updates[variable] = new_value if parameters is None: parameters = collect_differentiable_params(loss) if inputs is None: - inputs = collect_data_inputs([loss, *extra_outputs]) + inputs = collect_data_inputs([loss, *extra_outputs, *extra_updates.values()]) - updates: dict[SharedVariable, TensorVariable] = dict(rule(loss, parameters)) + updates: Updates = dict(rule(loss, parameters)) # Assigned per key rather than merged: SupportsKeysAndGetItem is invariant in its key type, so # dict.update rejects the narrower NonTrainableParameter keys. for parameter, new_value in collect_non_trainable_updates(loss).items(): updates[parameter] = new_value + # Folded in after the rule's and the model's, so a collision with either is caught rather than deciding + # by insertion order which of the two writes survives. + for variable, new_value in extra_updates.items(): + if variable in updates: + raise ValueError( + f"The extra update for {variable.name!r} writes a variable the training step already " + "writes, so the two writes cannot both take effect. Optimizer state and batch-norm " + "statistics are written by the step itself; drop this one, or fold what it does into the " + "expression that already writes the variable." + ) + updates[variable] = new_value + # Collected from the assembled updates rather than from the loss alone: a clock is read by a schedule # or a policy, which live in the updates, where an RNG or a running statistic is read by the model. for clock, next_count in collect_clock_updates( @@ -96,4 +132,4 @@ def compile_train( outputs = [loss, *extra_outputs] if extra_outputs else loss - return function(list(inputs), outputs, updates=updates, **(compile_kwargs or {})) + return function(list(inputs), outputs, updates=updates, **compile_kwargs) diff --git a/tests/optim/test_training.py b/tests/optim/test_training.py index 60114e7..08c73ae 100644 --- a/tests/optim/test_training.py +++ b/tests/optim/test_training.py @@ -2,7 +2,7 @@ import pytensor.tensor as pt import pytest -from pytensor import config +from pytensor import config, shared from pytensor.gradient import disconnected_grad, zero_grad from sklearn.datasets import load_digits, make_regression from sklearn.preprocessing import MinMaxScaler, OneHotEncoder, StandardScaler @@ -10,9 +10,9 @@ from pytensor_ml.activations import LeakyReLU from pytensor_ml.layers import BatchNorm2D, Linear, Sequential from pytensor_ml.loss import CrossEntropy, SquaredError, supervised_loss -from pytensor_ml.optim import adam, adamw, compile_train, sgd +from pytensor_ml.optim import adam, adamw, compile_train, cosine_schedule, sgd from pytensor_ml.optim.base import state_for -from pytensor_ml.params import trainable +from pytensor_ml.params import step_counter, trainable from pytensor_ml.pytensorf import collect_non_trainable_params, collect_trainable_params from pytensor_ml.state import initialize_params from pytensor_ml.util import DataLoader @@ -307,3 +307,169 @@ def test_compile_train_leaves_a_zero_grad_parameter_untouched_under_weight_decay for parameter, value in before.items() } assert moved == {"live_W": True, "live_b": True, "frozen_W": False, "frozen_b": False} + + +def test_extra_updates_write_state_no_gradient_produces(): + # The DQN shape from the other side: a target network kept as a Polyak average of the online weights. + # No gradient produces that write, so without extra_updates it cannot ride along in the training step. + X = pt.tensor("X", shape=(None, 4)) + online_layer = Linear("online", n_in=4, n_out=2) + prediction = online_layer(X) + parameters = collect_trainable_params(prediction) + initialize(parameters) + target_weight = shared(online_layer.W.get_value().copy(), name="target_W") + loss, target = supervised_loss(prediction, SquaredError(), ndim_out=2) + + step = compile_train( + loss, + sgd(1e-1), + parameters=parameters, + inputs=[X, target], + extra_updates={target_weight: 0.5 * target_weight + 0.5 * online_layer.W}, + ) + + online_before = online_layer.W.get_value().copy() + target_before = target_weight.get_value().copy() + + rng = np.random.default_rng(0) + features = rng.normal(size=(16, 4)).astype(config.floatX) + targets = np.ones((16, 2), dtype=config.floatX) + step(features, targets) + + # Updates are computed from the pre-update values, so the target lands halfway between where the two + # started -- neither staying put nor jumping to where the online weights ended up. + assert not np.allclose(online_layer.W.get_value(), online_before) + np.testing.assert_allclose( + target_weight.get_value(), 0.5 * target_before + 0.5 * online_before, rtol=1e-5 + ) + + +def test_extra_updates_reject_a_write_the_rule_already_makes(): + # Silently overwriting an optimizer buffer would leave the rule configured but not working, for the whole + # run, so the collision has to be loud. + p = trainable(np.array([2.0]), name="w") + loss = 0.5 * (p**2).sum() + rule = adam(1e-1) + first_moment = next(key for key in rule(loss, [p]) if key.name == "w/adam/first_moment") + + with pytest.raises(ValueError, match="already writes"): + compile_train(loss, rule, extra_updates={first_moment: first_moment * 0.0}, inputs=[]) + + +def test_extra_updates_reject_a_write_the_model_already_makes(): + # Batch-norm statistics are written by the model rather than the rule, and collide just the same. + X = pt.tensor("X", shape=(None, 4)) + prediction = Sequential(Linear("fc", n_in=4, n_out=4), BatchNorm2D("bn", n_in=4))(X) + parameters = collect_trainable_params(prediction) + initialize(parameters) + loss, target = supervised_loss(prediction, SquaredError(), ndim_out=2) + running_mean = next( + p for p in collect_non_trainable_params(prediction) if "running_mean" in p.name + ) + + with pytest.raises(ValueError, match="already writes"): + compile_train( + loss, + sgd(1e-2), + parameters=parameters, + inputs=[X, target], + extra_updates={running_mean: running_mean * 0.0}, + ) + + +def test_extra_updates_contribute_their_data_inputs(): + # An extra update may read data the loss never touches -- replay priorities, an importance weight. Unless + # those inputs are collected too, compiling raises MissingInputError. + p = trainable(np.array([2.0]), name="w") + priorities = shared(np.zeros(3, dtype=config.floatX), name="priorities") + fresh_priorities = pt.vector("fresh_priorities", shape=(3,)) + loss = 0.5 * (p**2).sum() + + step = compile_train(loss, sgd(1e-1), extra_updates={priorities: fresh_priorities}) + + step(fresh_priorities=np.array([1.0, 2.0, 3.0], dtype=config.floatX)) + np.testing.assert_allclose(priorities.get_value(), [1.0, 2.0, 3.0]) + + +def test_an_extra_update_reads_the_clock_and_advances_it_once(): + # An extra update is part of the step, so a clock it reads is a clock the step reads: it advances once + # per step, and the expression sees the count from before the advance, like the rule's updates do. + p = trainable(np.array([2.0]), name="w") + decayed = shared(np.array(1.0, dtype=config.floatX), name="decayed") + clock = step_counter(name="training_step") + schedule = cosine_schedule(1.0, 10) + loss = 0.5 * (p**2).sum() + + step = compile_train(loss, sgd(1e-1), extra_updates={decayed: schedule(clock)}, inputs=[]) + for _ in range(3): + step() + + assert int(clock.get_value()) == 3 + expected = float(schedule(pt.as_tensor(2, dtype="int64")).eval()) + np.testing.assert_allclose(decayed.get_value(), expected, rtol=1e-6) + + +def test_an_extra_update_that_draws_noise_advances_its_generator(): + # A perturbed step -- SGLD, exploration noise -- draws inside the update rather than the loss. Nothing + # else reads that generator, so unless the step advances it every call adds the identical perturbation. + p = trainable(np.array([2.0]), name="w") + perturbed = shared(np.zeros(3, dtype=config.floatX), name="perturbed") + noise_rng = shared(np.random.default_rng(0), name="noise_rng") + _, noise = pt.random.normal(size=(3,), rng=noise_rng, return_next_rng=True) + loss = 0.5 * (p**2).sum() + + step = compile_train( + loss, sgd(1e-1), extra_updates={perturbed: noise.astype(config.floatX)}, inputs=[] + ) + step() + first = perturbed.get_value().copy() + step() + + assert not np.allclose(first, perturbed.get_value()) + + +def test_updates_in_compile_kwargs_are_taken_as_extra_updates(): + # `updates` is what pytensor calls this, so it is the first place a caller looks. Forwarding it would + # collide with the compiler's own updates argument and raise a TypeError naming an internal function. + p = trainable(np.array([2.0]), name="w") + call_count = shared(np.array(0.0, dtype=config.floatX), name="call_count") + loss = 0.5 * (p**2).sum() + + step = compile_train( + loss, sgd(1e-1), inputs=[], compile_kwargs={"updates": {call_count: call_count + 1.0}} + ) + step() + step() + + assert float(call_count.get_value()) == 2.0 + + +def test_an_update_given_in_both_places_is_rejected(): + p = trainable(np.array([2.0]), name="w") + call_count = shared(np.array(0.0, dtype=config.floatX), name="call_count") + loss = 0.5 * (p**2).sum() + + with pytest.raises(ValueError, match="given twice"): + compile_train( + loss, + sgd(1e-1), + inputs=[], + extra_updates={call_count: call_count + 1.0}, + compile_kwargs={"updates": {call_count: call_count + 2.0}}, + ) + + +def test_compile_kwargs_is_not_mutated(): + # Taking `updates` out of the caller's dict would empty it, so compiling twice from one settings dict + # would silently drop the update the second time. + p = trainable(np.array([2.0]), name="w") + call_count = shared(np.array(0.0, dtype=config.floatX), name="call_count") + loss = 0.5 * (p**2).sum() + compile_kwargs = {"updates": {call_count: call_count + 1.0}} + + compile_train(loss, sgd(1e-1), inputs=[], compile_kwargs=compile_kwargs) + second = compile_train(loss, sgd(1e-1), inputs=[], compile_kwargs=compile_kwargs) + second() + + assert "updates" in compile_kwargs + assert float(call_count.get_value()) == 1.0