diff --git a/pytensor_ml/layers/dropout.py b/pytensor_ml/layers/dropout.py index 38c0281..0b176e7 100644 --- a/pytensor_ml/layers/dropout.py +++ b/pytensor_ml/layers/dropout.py @@ -6,6 +6,7 @@ from pytensor import config from pytensor.compile.sharedvalue import shared +from pytensor.tensor.random.variable import RandomGeneratorSharedVariable from pytensor_ml.base import Layer, UnaryLayerOp @@ -18,17 +19,53 @@ def build_inner_graph(self, X, mask): class Dropout(Layer): + """ + Zero a random fraction of its input, rescaling the rest so the expected sum is unchanged. + + Parameters + ---------- + name : str, optional + Name of the layer, used to name its output and its generators. Default "Dropout". + p : float + Probability of zeroing each element. Default 0.5. + random_state : int, Generator, or other seed, optional + Seed for the layer's masks. Draws are reproducible under a given seed, including when the layer is + applied at several points in a network. Seeded from fresh entropy when omitted. + + Attributes + ---------- + generators : list of RandomGeneratorSharedVariable + One generator per application of the layer, in the order they were applied. Set a value on these to + steer or restore the masks. + """ + def __init__(self, name: str | None = None, p: float = 0.5, random_state: Any | None = None): if p < 0.0 or p > 1.0: raise ValueError(f"Dropout probability has to be between 0 and 1, but got {p}") self.name = name if name else "Dropout" self.p = p - self.rng = shared(np.random.default_rng(random_state)) + self.generators: list[RandomGeneratorSharedVariable] = [] + self._generator_source = np.random.default_rng(random_state) + + def _own_generator(self) -> RandomGeneratorSharedVariable: + """Return a generator for one application of this layer, spawned from its seed. + + Each application draws its own mask, so each needs a generator of its own: a generator read by two + draws has no single next state, and nothing can advance it. Spawning keeps the whole layer + reproducible under ``random_state`` however many times it is applied. + """ + generator = shared( + self._generator_source.spawn(1)[0], name=f"{self.name}/rng_{len(self.generators)}" + ) + self.generators.append(generator) + return generator def __call__(self, X: pt.TensorLike) -> pt.TensorVariable: X = pt.as_tensor(X) p = pt.as_tensor(self.p, dtype=config.floatX) - _, mask = ptr.bernoulli(p=1 - p, size=X.shape, rng=self.rng, return_next_rng=True) + _, mask = ptr.bernoulli( + p=1 - p, size=X.shape, rng=self._own_generator(), return_next_rng=True + ) mask = mask.astype(config.floatX) X_masked = DropoutLayer( diff --git a/pytensor_ml/optim/base.py b/pytensor_ml/optim/base.py index 2896c24..f6193f3 100644 --- a/pytensor_ml/optim/base.py +++ b/pytensor_ml/optim/base.py @@ -57,11 +57,6 @@ def get_gradients( ------- list of TensorVariable One gradient per parameter, in the order of ``parameters``. - - Raises - ------ - DisconnectedInputError - If the loss has no gradient with respect to one of ``parameters``, naming the ones it cannot reach. """ if isinstance(loss_or_gradients, list | tuple): gradients = list(loss_or_gradients) diff --git a/pytensor_ml/pytensorf/collect.py b/pytensor_ml/pytensorf/collect.py index 7ba446a..c2d6ba5 100644 --- a/pytensor_ml/pytensorf/collect.py +++ b/pytensor_ml/pytensorf/collect.py @@ -110,12 +110,6 @@ def collect_clock_updates( ------- clock_updates : dict Mapping from each clock the graph reads to the expression for its next value. - - Raises - ------ - ValueError - If two clocks the graph reads hold different step counts. They all count training steps, so a - disagreement means some of them were restored from a checkpoint and others were not. """ counters = collect_step_counters(outputs) step_counts = {int(counter.get_value()) for counter in counters} diff --git a/pytensor_ml/pytensorf/compile.py b/pytensor_ml/pytensorf/compile.py index 03cb97b..8be9c26 100644 --- a/pytensor_ml/pytensorf/compile.py +++ b/pytensor_ml/pytensorf/compile.py @@ -1,10 +1,11 @@ +import warnings + from collections.abc import Sequence -from typing import cast import pytensor from pytensor import Mode -from pytensor.compile import Function, SharedVariable, get_mode +from pytensor.compile import Function, get_mode from pytensor.tensor.variable import Variable from pytensor_ml.pytensorf.collect import collect_graph_inputs @@ -13,6 +14,8 @@ SeedSequenceSeed, atleast_list, collect_default_updates, + find_generators_drawn_from, + find_rng_nodes, reseed_rngs, ) @@ -37,8 +40,9 @@ def function( outputs : Variable or list of Variable Outputs of the compiled function. random_seed : int, array-like of int, or SeedSequence, optional - Seed used to reseed the graph's shared generators. They are replaced whether or not a seed is - given, so omitting it reseeds from fresh entropy rather than leaving them untouched. + Seed used to replace the graph's shared generators, making the compiled function's draws + reproducible. The generators are left as they are when omitted, so compiling has no effect on a + generator the caller seeded, or on one another function is drawing from. mode : Mode or str, optional PyTensor mode used to compile the function. **kwargs @@ -49,16 +53,39 @@ def function( Function The compiled function. """ - rng_updates = collect_default_updates( - inputs=[inp.variable if isinstance(inp, pytensor.In) else inp for inp in inputs], - outputs=[ - out.variable if isinstance(out, pytensor.Out) else out for out in atleast_list(outputs) - ], - ) - - if rng_updates: - rngs = cast(list[SharedVariable], list(rng_updates)) - reseed_rngs(rngs, random_seed) + updates = dict(kwargs.pop("updates", {})) + input_variables = [inp.variable if isinstance(inp, pytensor.In) else inp for inp in inputs] + # Updates count as readers: a generator a rule draws its noise from is read by the update expression + # and often by nothing else, and one read by both an output and an update is read twice. + read_variables = [ + *(out.variable if isinstance(out, pytensor.Out) else out for out in atleast_list(outputs)), + *updates.values(), + ] + + if random_seed is not None: + reseed_rngs(find_rng_nodes(read_variables), random_seed) + + with warnings.catch_warnings(): + # This warns for a generator with several distinct draws and returns no update for it. The check + # below reports that better, and only once the caller's own updates are known. + warnings.filterwarnings( + "ignore", message="RNG Variable .* multiple distinct clients", category=UserWarning + ) + rng_updates = collect_default_updates(inputs=input_variables, outputs=read_variables) + + frozen = [ + generator + for generator in find_generators_drawn_from(read_variables) + if generator not in rng_updates and generator not in updates + ] + if frozen: + raise ValueError( + f"The graph draws from {[str(generator.name or generator) for generator in frozen]}, which " + "nothing advances, so every call would repeat the same values. A generator read by two " + "different draws has no single next state, which is the usual cause. Give each draw its own " + "generator, thread one through with `next_rng, draw = pt.random.normal(rng=rng, " + "return_next_rng=True)`, or pass an update for it yourself." + ) base_mode = get_mode(mode) mode = Mode( @@ -69,7 +96,7 @@ def function( return pytensor.function( inputs, outputs, - updates={**rng_updates, **kwargs.pop("updates", {})}, + updates={**rng_updates, **updates}, mode=mode, **kwargs, ) diff --git a/pytensor_ml/pytensorf/rng.py b/pytensor_ml/pytensorf/rng.py index 7778238..54e29a4 100644 --- a/pytensor_ml/pytensorf/rng.py +++ b/pytensor_ml/pytensorf/rng.py @@ -46,6 +46,34 @@ def reseed_rngs(rngs: Sequence[SharedVariable], seed: SeedSequenceSeed) -> None: rng.set_value(np.random.Generator(bit_generator), borrow=True) +def find_generators_drawn_from( + outputs: Sequence[Variable], +) -> list[RandomGeneratorSharedVariable]: + """ + Return the shared generators the graph draws from, excluding any it only hands back. + + A generator returned as an output is read without being consumed, so it needs no update; one an op draws + from does, and a graph that draws from a generator nothing advances repeats the same values forever. + + Parameters + ---------- + outputs : sequence of Variable + Graph outputs to trace back from. + + Returns + ------- + list of RandomGeneratorSharedVariable + The generators some node in the graph draws from, in graph-input order. + """ + fgraph = FunctionGraph(outputs=list(outputs), clone=False) + return [ + generator + for generator in fgraph.inputs + if isinstance(generator, RandomGeneratorSharedVariable) + and any(not isinstance(client.op, Output) for client, _ in fgraph.clients[generator]) + ] + + def collect_default_updates_inner_fgraph(node: Apply) -> dict[Variable, Variable]: """Collect default RNG updates from a node carrying an inner function graph, mapped to outer variables.""" op = node.op diff --git a/tests/test_pytensorf.py b/tests/test_pytensorf.py index 251f311..bcc6357 100644 --- a/tests/test_pytensorf.py +++ b/tests/test_pytensorf.py @@ -11,11 +11,13 @@ zero_grad, ) from pytensor.graph.traversal import ancestors +from pytensor.tensor import random as ptr from pytensor_ml.layers import Dropout, DropoutLayer, Linear, Sequential from pytensor_ml.pytensorf import ( collect_trainable_params, compile_predict, + function, rewrite_for_prediction, rewrite_pregrad, ) @@ -90,3 +92,145 @@ def test_rewrite_pregrad_leaves_a_fully_detached_parameter_disconnected(): with pytest.raises(DisconnectedInputError): grad(rewrite_pregrad(loss), W) + + +def test_compiling_leaves_a_generator_another_function_is_drawing_from_alone(): + """Compiling used to replace every generator it touched, so building a second function mid-training + jumped the first one's noise stream.""" + + def four_draws(with_intervening_compile): + rng = shared(np.random.default_rng(0), name="rng") + _, noise = ptr.normal(rng=rng, return_next_rng=True) + draw = function([], noise) + drawn = [float(draw()) for _ in range(2)] + if with_intervening_compile: + _, other_noise = ptr.normal(rng=rng, return_next_rng=True) + function([], other_noise) + return drawn + [float(draw()) for _ in range(2)] + + assert four_draws(with_intervening_compile=True) == four_draws(with_intervening_compile=False) + + +@pytest.mark.parametrize("applications", [1, 2], ids=["applied_once", "applied_twice"]) +def test_a_seeded_dropout_reproduces_across_identical_runs(applications): + """The seed a caller puts on a layer has to survive compilation, or `random_state=` means nothing. It has + to hold however many times the layer is applied, since each application draws off its own generator.""" + + def dropout_masks(): + X = pt.tensor("X", shape=(None, 3)) + dropout = Dropout(p=0.5, random_state=0) + layers = [Linear("fc", n_in=3, n_out=3)] + for _ in range(applications): + layers.append(dropout) + prediction = Sequential(*layers)(X) + for parameter in collect_trainable_params(prediction): + parameter.set_value(np.ones_like(parameter.get_value())) + forward = function([X], prediction) + features = np.ones((4, 3), dtype=config.floatX) + return [forward(features) for _ in range(3)] + + for first, second in zip(dropout_masks(), dropout_masks()): + np.testing.assert_allclose(first, second) + + +def test_random_seed_makes_an_unseeded_graph_reproducible(): + def two_draws(): + rng = shared(np.random.default_rng(), name="rng") # deliberately unseeded + _, noise = ptr.normal(rng=rng, return_next_rng=True) + draw = function([], noise, random_seed=42) + return [float(draw()) for _ in range(2)] + + assert two_draws() == two_draws() + + +def test_a_reused_dropout_instance_keeps_drawing_new_masks(): + """Using one Dropout object at two points in a network is an ordinary thing to write, and it used to + freeze the mask for the whole run.""" + X = pt.tensor("X", shape=(None, 3)) + dropout = Dropout(p=0.5, random_state=0) + prediction = Sequential( + Linear("fc1", n_in=3, n_out=4), dropout, Linear("fc2", n_in=4, n_out=2), dropout + )(X) + for parameter in collect_trainable_params(prediction): + parameter.set_value(np.ones_like(parameter.get_value())) + + forward = function([X], prediction) + features = np.ones((4, 3), dtype=config.floatX) + outputs = [forward(features) for _ in range(4)] + + assert ( + len(dropout.generators) == 2 + ) # one per application, so neither draw is starved of updates + assert all( + np.any(outputs[i] != outputs[j]) for i in range(4) for j in range(i + 1, 4) + ) # a fresh mask on every call, not just eventually + + +def test_a_returned_generator_is_the_one_the_caller_asked_for(): + """Returning a generator alongside a draw from it is not two draws: an output reads the generator without + consuming it, so it must keep reading the caller's own.""" + rng = shared(np.random.default_rng(0), name="rng") + _, noise = ptr.normal(rng=rng, return_next_rng=True) + draw = function([], [noise, rng]) + + expected = rng.get_value(borrow=True).bit_generator.state + returned = draw()[1].bit_generator.state + + assert returned == expected + + +def test_a_generator_read_only_by_an_update_still_advances(): + """A rule that adds noise to its step reads a generator the outputs never touch, so collecting updates + from the outputs alone left that generator frozen and every step took the identical perturbation.""" + rng = shared(np.random.default_rng(0), name="rng") + parameter = shared(np.zeros(()), name="w") + _, noise = ptr.normal(rng=rng, return_next_rng=True) + + step = function([], parameter**2, updates={parameter: parameter + noise}) + + increments = [] + for _ in range(3): + before = float(parameter.get_value()) + step() + increments.append(float(parameter.get_value()) - before) + + assert not np.allclose(increments[0], increments[1]) + assert not np.allclose(increments[1], increments[2]) + + +def test_two_distinct_draws_off_one_generator_are_rejected(): + """One generator cannot serve two different draws: it has no single next state, so nothing can advance + it and every call would repeat. pytensor calls such a graph inconsistent and threads no update, which + would otherwise show up as a training run whose noise never changes.""" + rng = shared(np.random.default_rng(0), name="rng") + _, normal_draw = ptr.normal(rng=rng, return_next_rng=True) + _, uniform_draw = ptr.uniform(rng=rng, return_next_rng=True) + + with pytest.raises(ValueError, match="which nothing advances"): + function([], [normal_draw, uniform_draw]) + + +def test_a_shared_generator_is_accepted_when_the_caller_advances_it(): + """The check is about the assembled updates, not about how the graph looks: a caller who threads the + generator themselves has answered the question, and their update stands.""" + rng = shared(np.random.default_rng(0), name="rng") + next_rng, normal_draw = ptr.normal(rng=rng, return_next_rng=True) + _, uniform_draw = ptr.uniform(rng=next_rng, return_next_rng=True) + + draw = function([], [normal_draw, uniform_draw], updates={rng: next_rng}) + + first, second = draw(), draw() + assert float(first[0]) != float(second[0]) # the draw off the generator advances + assert float(first[1]) != float(second[1]) # and so does the one off the threaded next state + + +def test_a_draw_shared_between_an_output_and_an_update_is_rejected(): + """An update expression is a reader like any other, so a generator drawn from by both an output and an + update is read twice and cannot be advanced once.""" + rng = shared(np.random.default_rng(0), name="rng") + accumulator = shared(np.zeros(()), name="accumulator") + _, output_draw = ptr.normal(rng=rng, return_next_rng=True) + _, update_draw = ptr.uniform(rng=rng, return_next_rng=True) + + with pytest.raises(ValueError, match="which nothing advances"): + function([], output_draw, updates={accumulator: accumulator + update_draw})