diff --git a/pytensor_ml/model.py b/pytensor_ml/model.py index fdae27c..29004d5 100644 --- a/pytensor_ml/model.py +++ b/pytensor_ml/model.py @@ -8,7 +8,7 @@ from pytensor_ml.loss import Loss, supervised_loss from pytensor_ml.params import TrainableParameter, collect_trainable_params from pytensor_ml.pytensorf import compile_predict -from pytensor_ml.state import InitializationScheme, initialize_params +from pytensor_ml.state import InitializationSchemeLike, initialize_params class Model: @@ -26,7 +26,7 @@ def weights(self) -> list[TrainableParameter]: def initialize( self, - scheme: InitializationScheme = "xavier_normal", + scheme: InitializationSchemeLike = "xavier_normal", seed: int | np.random.Generator | None = None, ) -> "Model": """ @@ -34,8 +34,9 @@ def initialize( Parameters ---------- - scheme : str - Initialization scheme for the weights. Default 'xavier_normal'. + scheme : str or Initializer + Initialization scheme for the weights: the name of a built-in scheme, or an + :class:`~pytensor_ml.state.Initializer` instance. Default 'xavier_normal'. seed : int or numpy Generator, optional Seed for reproducible initialization. """ @@ -84,8 +85,7 @@ def predict(self, X_values: np.ndarray) -> np.ndarray: self.y, inputs=[self.X], compile_kwargs=self._compile_kwargs ) - result = self._predict_fn(X_values) - return result if isinstance(result, np.ndarray) else np.asarray(result) + return np.asarray(self._predict_fn(X_values)) def __str__(self): return debugprint(self.y, file="str") diff --git a/pytensor_ml/state.py b/pytensor_ml/state.py index 98cc09c..42e9c0d 100644 --- a/pytensor_ml/state.py +++ b/pytensor_ml/state.py @@ -19,20 +19,11 @@ class Initializer(ABC): """ Base class for parameter initializers. - Can be used in two ways: - - As a class: `XavierNormalInitializer(param, rng)` - directly initializes - - As an instance: `init = XavierNormalInitializer(); init(param, rng)` + Subclasses implement :meth:`sample`. Calling an instance assigns a freshly sampled value to a + parameter in place, while :func:`initialize_params` calls :meth:`sample` directly and leaves the + assignment to its caller. """ - def __new__(cls, param: SharedVariable | None = None, rng: RandomState | None = None): - # If called with a param, act as a function and initialize directly - if param is not None: - instance = object.__new__(cls) - cls.__init__(instance) - return instance(param, rng) - # Otherwise, return an instance for later use - return object.__new__(cls) - def __call__(self, param: SharedVariable, rng: RandomState | None = None) -> SharedVariable: param.set_value(self._sample_like(param, rng)) return param @@ -60,29 +51,25 @@ def sample(self, shape: tuple[int, ...], dtype: str, rng: np.random.Generator) - class XavierUniformInitializer(Initializer): def sample(self, shape: tuple[int, ...], dtype: str, rng: np.random.Generator) -> np.ndarray: - scale = np.sqrt(6.0 / np.sum([x for x in shape if x is not None])) + scale = np.sqrt(6.0 / np.sum(shape)) return rng.uniform(-scale, scale, size=shape).astype(dtype) class XavierNormalInitializer(Initializer): def sample(self, shape: tuple[int, ...], dtype: str, rng: np.random.Generator) -> np.ndarray: - scale = np.sqrt(2.0 / np.sum([x for x in shape if x is not None])) + scale = np.sqrt(2.0 / np.sum(shape)) return rng.normal(0, scale, size=shape).astype(dtype) class CustomInitializer(Initializer): - def __new__( - cls, - sample_fn: SamplingFunction | None = None, - param: SharedVariable | None = None, - rng: RandomState | None = None, - ): - instance = object.__new__(cls) - if sample_fn is not None: - instance._sample_fn = sample_fn - if param is not None: - return instance(param, rng) - return instance + """ + Initializer built from a sampling function. + + Parameters + ---------- + sample_fn : callable + ``(shape, dtype, rng) -> ndarray``, returning the initial value for one parameter. + """ def __init__(self, sample_fn: SamplingFunction): self._sample_fn = sample_fn @@ -91,17 +78,19 @@ def sample(self, shape: tuple[int, ...], dtype: str, rng: np.random.Generator) - return self._sample_fn(shape, dtype, rng) -_INIT_FUNCTIONS: dict[str, type[Initializer]] = { +_INITIALIZERS: dict[str, type[Initializer]] = { "zeros": ZeroInitializer, "xavier_uniform": XavierUniformInitializer, "xavier_normal": XavierNormalInitializer, "unit_uniform": UnitUniformInitializer, } +InitializationSchemeLike = InitializationScheme | Initializer + def initialize_params( params: Sequence[SharedVariable], - scheme: InitializationScheme = "xavier_normal", + scheme: InitializationSchemeLike = "xavier_normal", rng: RandomState | None = None, ) -> list[np.ndarray]: """ @@ -112,7 +101,8 @@ def initialize_params( params SharedVariables to initialize values for. scheme - Initialization scheme to use. + Initialization scheme to use: the name of a built-in scheme, or any :class:`Initializer` + instance (including a :class:`CustomInitializer` wrapping your own sampling function). rng Random number generator. If None, a new one is created. @@ -121,11 +111,8 @@ def initialize_params( list of np.ndarray Initialized values matching the shapes and dtypes of params. """ + # Resolve once and share: a seed handed to each _sample_like call would repeat draws across parameters. rng = np.random.default_rng(rng) - initializer = _INIT_FUNCTIONS[scheme]() - results = [] - for var in params: - value = var.get_value() - results.append(initializer.sample(value.shape, str(value.dtype), rng)) - return results + initializer = scheme if isinstance(scheme, Initializer) else _INITIALIZERS[scheme]() + return [initializer._sample_like(param, rng) for param in params] diff --git a/tests/rewriting/test_layers.py b/tests/rewriting/test_layers.py index 72b8ce4..8df8ad9 100644 --- a/tests/rewriting/test_layers.py +++ b/tests/rewriting/test_layers.py @@ -44,8 +44,16 @@ def test_remove_dropout(feature_extractor_and_rng): assert len([node.op for node in fg.apply_nodes if isinstance(node.op, DropoutLayer)]) == 0 -def test_rewrite_batch_stats_to_running_average_stats(): - feature_extractor = Sequential(Linear("Layer_1", n_in=6, n_out=3), BatchNorm2D()) +@pytest.mark.parametrize( + "consumed_downstream", [False, True], ids=["as_graph_output", "consumed_by_a_layer"] +) +def test_rewrite_batch_stats_to_running_average_stats(consumed_downstream): + layers = [Linear("Layer_1", n_in=6, n_out=3), BatchNorm2D()] + if consumed_downstream: + # A downstream SymbolicOp type-checks its inputs more strictly than a graph output does. + layers.append(Linear("Layer_2", n_in=3, n_out=1)) + + feature_extractor = Sequential(*layers) X = pt.tensor("X", shape=(None, 6)) latent = feature_extractor(X) diff --git a/tests/test_model.py b/tests/test_model.py index 22598f0..13e78a1 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -3,6 +3,8 @@ from pytensor import config +import pytensor_ml.model + from pytensor_ml.layers import BatchNorm2D, Linear, Sequential from pytensor_ml.loss import SquaredError from pytensor_ml.model import Model @@ -10,107 +12,76 @@ class TestModelPredict: - def test_simple_network(self): + def test_matches_a_hand_computed_forward_pass(self): X = pt.tensor("X", shape=(None, 6)) - mlp = Sequential(Linear("fc1", n_in=6, n_out=3), Linear("fc2", n_in=3, n_out=1)) - y = mlp(X) + fc1 = Linear("fc1", n_in=6, n_out=3) + fc2 = Linear("fc2", n_in=3, n_out=1) + model = Model(X, Sequential(fc1, fc2)(X)).initialize(seed=42) - model = Model(X, y) - model.initialize(seed=42) + X_test = np.random.default_rng(0).normal(size=(10, 6)).astype(config.floatX) + hidden = X_test @ fc1.W.get_value() + fc1.b.get_value() + expected = hidden @ fc2.W.get_value() + fc2.b.get_value() - X_test = np.random.randn(10, 6).astype(config.floatX) result = model.predict(X_test) assert result.shape == (10, 1) assert result.dtype == config.floatX + np.testing.assert_allclose(result, expected, rtol=1e-5) - def test_with_batchnorm(self): - X = pt.tensor("X", shape=(None, 6)) - network = Sequential( - Linear("fc1", n_in=6, n_out=3), - BatchNorm2D("bn1", n_in=3), - Linear("fc2", n_in=3, n_out=1), - ) - y = network(X) - - model = Model(X, y) - model.initialize(seed=42) - - X_test = np.random.randn(10, 6).astype(config.floatX) - result = model.predict(X_test) - - assert result.shape == (10, 1) - - def test_predict_uses_running_stats(self): - """Verify that predict uses running stats (via rewrite) not batch stats.""" + def test_normalizes_with_running_stats_not_batch_stats(self): X = pt.tensor("X", shape=(None, 4)) fc1 = Linear("fc1", n_in=4, n_out=4) bn = BatchNorm2D("bn1", n_in=4) - network = Sequential(fc1, bn) - y = network(X) - - model = Model(X, y) - model.initialize(seed=42) + model = Model(X, Sequential(fc1, bn)(X)).initialize(seed=42) - # Set specific running stats bn.running_mean.set_value(np.array([1.0, 2.0, 3.0, 4.0], dtype=config.floatX)) - bn.running_var.set_value(np.array([1.0, 1.0, 1.0, 1.0], dtype=config.floatX)) + bn.running_var.set_value(np.ones(4, dtype=config.floatX)) - # Predict with two different batches - should give same normalization - # if using running stats (batch stats would differ) - X1 = np.random.randn(5, 4).astype(config.floatX) - X2 = np.random.randn(20, 4).astype(config.floatX) - - # Get FC output before normalization for both batches - fc_weight = fc1.W.get_value() - fc_bias = fc1.b.get_value() - fc_out1 = X1 @ fc_weight + fc_bias - fc_out2 = X2 @ fc_weight + fc_bias - - # Expected: normalized using running stats - scale = bn.scale.get_value() - loc = bn.loc.get_value() - running_mean = bn.running_mean.get_value() - running_var = bn.running_var.get_value() - - expected1 = (fc_out1 - running_mean) / np.sqrt(running_var + bn.epsilon) * scale + loc - expected2 = (fc_out2 - running_mean) / np.sqrt(running_var + bn.epsilon) * scale + loc + rng = np.random.default_rng(0) + # Two batch sizes: batch statistics would differ between them, running statistics cannot. + for n_rows in (5, 20): + X_test = rng.normal(size=(n_rows, 4)).astype(config.floatX) + fc_out = X_test @ fc1.W.get_value() + fc1.b.get_value() + standardized = (fc_out - bn.running_mean.get_value()) / np.sqrt( + bn.running_var.get_value() + bn.epsilon + ) + expected = standardized * bn.scale.get_value() + bn.loc.get_value() + + np.testing.assert_allclose(model.predict(X_test), expected, rtol=1e-5) + + def test_compiles_once_and_reuses_the_function(self, monkeypatch): + X = pt.tensor("X", shape=(None, 4)) + model = Model(X, Linear("fc1", n_in=4, n_out=2)(X)).initialize(seed=42) - result1 = model.predict(X1) - result2 = model.predict(X2) + uncounted_compile_predict = pytensor_ml.model.compile_predict + compile_count = 0 - np.testing.assert_allclose(result1, expected1, rtol=1e-5) - np.testing.assert_allclose(result2, expected2, rtol=1e-5) + def counting_compile_predict(*args, **kwargs): + nonlocal compile_count + compile_count += 1 + return uncounted_compile_predict(*args, **kwargs) - def test_compile_train_reduces_loss(self): - X = pt.tensor("X", shape=(None, 4)) - y = Sequential(Linear("fc1", n_in=4, n_out=8), Linear("fc2", n_in=8, n_out=1))(X) - model = Model(X, y).initialize(seed=0) + monkeypatch.setattr(pytensor_ml.model, "compile_predict", counting_compile_predict) - step = model.compile_train(sgd(learning_rate=1e-2), SquaredError(), ndim_out=2) + X_test = np.random.default_rng(0).normal(size=(5, 4)).astype(config.floatX) + first = model.predict(X_test) + second = model.predict(X_test) - rng = np.random.default_rng(0) - X_batch = rng.normal(size=(64, 4)).astype(config.floatX) - target = rng.normal(size=(64, 1)).astype(config.floatX) - history = [float(step(X_batch, target)) for _ in range(50)] - assert history[-1] < history[0] + assert compile_count == 1 + np.testing.assert_array_equal(first, second) - def test_predict_caches_function(self): - """Verify that predict function is compiled once and reused.""" - X = pt.tensor("X", shape=(None, 4)) - y = Linear("fc1", n_in=4, n_out=2)(X) - model = Model(X, y) - model.initialize(seed=42) +def test_compile_train_reduces_loss(): + X = pt.tensor("X", shape=(None, 4)) + y = Sequential(Linear("fc1", n_in=4, n_out=8), Linear("fc2", n_in=8, n_out=1))(X) + model = Model(X, y).initialize(seed=0) - assert model._predict_fn is None + step = model.compile_train(sgd(learning_rate=1e-2), SquaredError(), ndim_out=2) - X_test = np.random.randn(5, 4).astype(config.floatX) - model.predict(X_test) + rng = np.random.default_rng(0) + X_batch = rng.normal(size=(64, 4)).astype(config.floatX) + target = rng.normal(size=(64, 1)).astype(config.floatX) - assert model._predict_fn is not None - fn_id = id(model._predict_fn) + history = [float(step(X_batch, target)) for _ in range(50)] - # Second call should reuse same function - model.predict(X_test) - assert id(model._predict_fn) == fn_id + assert history[-1] < history[0] diff --git a/tests/test_state.py b/tests/test_state.py index 959c00e..29a7e31 100644 --- a/tests/test_state.py +++ b/tests/test_state.py @@ -1,25 +1,54 @@ +from typing import get_args + import numpy as np +import pytest + +from pytensor_ml.params import collect_trainable_params, trainable +from pytensor_ml.state import ( + _INITIALIZERS, + CustomInitializer, + InitializationScheme, + ZeroInitializer, + initialize_params, +) + -from pytensor_ml.params import collect_trainable_params -from pytensor_ml.state import initialize_params +def test_scheme_names_match_the_initializer_registry(): + assert set(get_args(InitializationScheme)) == set(_INITIALIZERS) class TestInitializeParams: - def test_xavier_normal(self, simple_network): + @pytest.mark.parametrize("scheme", sorted(_INITIALIZERS)) + def test_values_match_parameter_shapes_and_dtypes(self, scheme, simple_network): X, y = simple_network params = collect_trainable_params(y) - values = initialize_params(params, scheme="xavier_normal", rng=np.random.default_rng(42)) + + values = initialize_params(params, scheme=scheme, rng=np.random.default_rng(0)) assert len(values) == len(params) for param, val in zip(params, values): assert val.shape == param.get_value().shape assert str(val.dtype) == str(param.get_value().dtype) - def test_zeros(self, simple_network): + def test_xavier_normal_scales_by_fan_in_plus_fan_out(self, simple_network): + X, y = simple_network + weight = next(p for p in collect_trainable_params(y) if p.name == "fc1_W") + fan_sum = sum(weight.get_value().shape) + + [value] = initialize_params([weight], scheme="xavier_normal", rng=np.random.default_rng(42)) + + assert value.std() == pytest.approx(np.sqrt(2.0 / fan_sum), rel=0.05) + # Xavier uniform targets this same variance; only a normal draw crosses its hard bound. + assert np.abs(value).max() > np.sqrt(6.0 / fan_sum) + + @pytest.mark.parametrize("scheme", ["zeros", ZeroInitializer()], ids=["by_name", "by_instance"]) + def test_zeros(self, scheme, simple_network): X, y = simple_network params = collect_trainable_params(y) - values = initialize_params(params, scheme="zeros") + values = initialize_params(params, scheme=scheme) + + assert len(values) == len(params) for val in values: np.testing.assert_array_equal(val, 0) @@ -32,3 +61,32 @@ def test_reproducible_with_seed(self, simple_network): for v1, v2 in zip(values1, values2): np.testing.assert_array_equal(v1, v2) + + def test_parameters_do_not_all_receive_the_same_draws(self): + # The seed must be an int: passing a Generator masks the regression this guards. + first = trainable(np.zeros((4, 4), dtype="float64"), "first") + second = trainable(np.zeros((4, 4), dtype="float64"), "second") + + values = initialize_params([first, second], scheme="unit_uniform", rng=0) + + assert not np.array_equal(values[0], values[1]) + + def test_accepts_a_custom_initializer(self, simple_network): + X, y = simple_network + params = collect_trainable_params(y) + constant = CustomInitializer(lambda shape, dtype, rng: np.full(shape, 7.0, dtype=dtype)) + + values = initialize_params(params, scheme=constant) + + assert len(values) == len(params) + for val in values: + np.testing.assert_array_equal(val, 7.0) + + +def test_calling_an_initializer_assigns_the_parameter_in_place(): + param = trainable(np.ones(3, dtype="float64"), "w") + + returned = ZeroInitializer()(param) + + np.testing.assert_array_equal(param.get_value(), 0) + assert returned is param