diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 0d98dce..95a647a 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -33,7 +33,9 @@ jobs: python -m pip install --upgrade pip setuptools wheel # Install runtime + test deps explicitly: editable metadata from # setup.py is not reliably applied by current pip/setuptools. - python -m pip install "numpy>=1.7" "scipy>=0.12" "GPy>=1.15" matplotlib h5py tables "coverage>=7.0" + python -m pip install "numpy>=1.7" "scipy>=0.12" matplotlib h5py tables "coverage>=7.0" + # CIP-0001 Phase 3 needs GPy psicomp chunk_size (merged in GPy #1299; not yet on PyPI). + python -m pip install "git+https://github.com/SheffieldML/GPy.git@main" python -m pip install -e . --no-deps - name: Run unit tests with coverage @@ -63,7 +65,8 @@ jobs: sudo apt-get update sudo apt-get install -y --no-install-recommends libopenmpi-dev openmpi-bin python -m pip install --upgrade pip setuptools wheel - python -m pip install "numpy>=1.7" "scipy>=0.12" "GPy>=1.15" matplotlib h5py tables mpi4py + python -m pip install "numpy>=1.7" "scipy>=0.12" matplotlib h5py tables mpi4py + python -m pip install "git+https://github.com/SheffieldML/GPy.git@main" python -m pip install -e . --no-deps - name: MPI parity (2 ranks vs serial) diff --git a/README.md b/README.md index 9a67c09..4a81139 100644 --- a/README.md +++ b/README.md @@ -21,6 +21,8 @@ To run the suite locally: python -m pip install -e ".[test]" # or, if editable extras are unreliable: python -m pip install "numpy>=1.7" "scipy>=0.12" "GPy>=1.15" matplotlib h5py tables "coverage>=7.0" +# For psi_chunk_size / psi_threads (CIP-0001 Phase 3) until the next GPy release: +# python -m pip install "git+https://github.com/SheffieldML/GPy.git@main" python -m pip install -e . --no-deps python -m coverage run -m unittest discover -s deepgp/testing -p 'test*.py' -v @@ -47,3 +49,15 @@ single-machine role of MPI. ``mpi_comm`` with ``back_constraint=False`` raises ``ValueError`` (that configuration produced silently wrong results). An optional GitHub Actions job runs a 2-rank parity check while MPI remains. + +Psi chunking (optional) +----------------------- + +Large-N models can bound memory and use multiple cores via GPy's chunked RBF / +linear psi-statistics (CIP-0001 Phase 3):: + + m = deepgp.DeepGP(..., psi_chunk_size=500, psi_threads=4) + +Defaults (``psi_chunk_size=None``, ``psi_threads=1``) preserve previous +behaviour. Threading runs inside GPy's pure psi helpers; PyDeepGP itself stays +serial at the kernel API. diff --git a/cip/cip0001.md b/cip/cip0001.md index 2014e9e..dccc43c 100644 --- a/cip/cip0001.md +++ b/cip/cip0001.md @@ -27,7 +27,7 @@ title: Replace Theano and MPI with modern NumPy - [x] Proposed - Initial idea documented (open for Discussion feedback) - [x] Accepted - Phase 1 started on branch `cip0001-numpy-encoder` (regularisation: keep no-op + FutureWarning) -- [x] In Progress - Phase 2 (MPI correctness / deprecation) under way +- [x] In Progress - Phase 3 (chunked psi-statistics) under way - [ ] Implemented - Work complete, awaiting verification - [ ] Closed - Verified and complete - [ ] Rejected - Will not be implemented @@ -256,8 +256,8 @@ pluggable reducer in R14, so the MPI-specific code shrinks to that one place. - [ ] Discussion opened - [x] Accept (partial): regularisation kept as no-op with `FutureWarning` (R2) - [x] Phase 1: NumPy encoder (R1–R7) — merged (PR #37) -- [x] Phase 2: MPI correctness and deprecation (R8–R11) on `cip0001-mpi-deprecate` — typemap removed; unsupported configs raise; FutureWarning; optional MPI CI job -- [ ] Phase 3: chunked psi-statistics (R12 in GPy, R13–R16) +- [x] Phase 2: MPI correctness and deprecation (R8–R11) — merged (PR #38) +- [x] Phase 3: chunked psi-statistics — GPy R12/R13 merged (GPy PR #1299); PyDeepGP R14–R16 on `cip0001-psi-chunking` (`psi_chunk_size` / `psi_threads`) - [ ] Many-core benchmarks recorded - [ ] Phase 4: MPI removed (R17–R24) or reduced to the reducer hook diff --git a/deepgp/inference/svi_vardtc.py b/deepgp/inference/svi_vardtc.py index 46a26a2..bc92392 100644 --- a/deepgp/inference/svi_vardtc.py +++ b/deepgp/inference/svi_vardtc.py @@ -26,9 +26,11 @@ class SVI_VarDTC(LatentFunctionInference): """ const_jitter = 1e-6 - def __init__(self, mpi_comm=None, mpi_root=0): + def __init__(self, mpi_comm=None, mpi_root=0, psi_chunk_size=None, psi_threads=1): self.mpi_comm = mpi_comm self.mpi_root=mpi_root + self.psi_chunk_size = psi_chunk_size + self.psi_threads = 1 if psi_threads is None else int(psi_threads) def get_trYYT(self, Y): return np.sum(np.square(Y)) @@ -41,10 +43,34 @@ def get_YYTfactor(self, Y): return jitchol(tdot(Y)) def gatherPsiStat(self, kern, X, Z, Y, beta, uncertain_inputs): + from ..util.psi_chunking import configure_kern_psicomp, row_slices assert beta.size == 1 - if uncertain_inputs: + psicomp_chunked = configure_kern_psicomp(kern, self.psi_chunk_size, self.psi_threads) + N = X.shape[0] + chunk = self.psi_chunk_size + use_serial_chunks = ( + uncertain_inputs + and chunk is not None + and int(chunk) > 0 + and int(chunk) < N + and not psicomp_chunked + ) + + if use_serial_chunks: + psi1_parts = [] + psi2 = None + psi0 = 0.0 + for sl in row_slices(N, chunk): + Xc = X[sl] + psi0 = psi0 + kern.psi0(Z, Xc).sum() + psi1_parts.append(kern.psi1(Z, Xc) * beta) + psi2_c = kern.psi2(Z, Xc) * beta + psi2 = psi2_c if psi2 is None else psi2 + psi2_c + psi1 = np.vstack(psi1_parts) + psi0 = psi0 * beta + elif uncertain_inputs: psi0 = kern.psi0(Z, X).sum()*beta psi1 = kern.psi1(Z, X)*beta psi2 = kern.psi2(Z, X)*beta diff --git a/deepgp/inference/vardtc.py b/deepgp/inference/vardtc.py index b9b9f5b..2005cd0 100644 --- a/deepgp/inference/vardtc.py +++ b/deepgp/inference/vardtc.py @@ -29,9 +29,11 @@ class VarDTC(LatentFunctionInference): """ const_jitter = 1e-6 - def __init__(self, mpi_comm=None): + def __init__(self, mpi_comm=None, psi_chunk_size=None, psi_threads=1): self.mpi_comm = mpi_comm + self.psi_chunk_size = psi_chunk_size + self.psi_threads = 1 if psi_threads is None else int(psi_threads) self.Y_speedup = False # Replace Y with the cholesky factor of YY.T, but the computation of posterior object will be skipped. def get_trYYT(self, Y): @@ -45,20 +47,46 @@ def get_YYTfactor(self, Y): return jitchol(tdot(Y)) def gatherPsiStat(self, kern, X, Z, Y, beta, uncertain_inputs): + from ..util.psi_chunking import configure_kern_psicomp, row_slices assert beta.size == 1 - if uncertain_inputs: + # Prefer GPy R12 chunking/threading when the kernel supports it. + psicomp_chunked = configure_kern_psicomp(kern, self.psi_chunk_size, self.psi_threads) + N = X.shape[0] + chunk = self.psi_chunk_size + use_serial_chunks = ( + uncertain_inputs + and chunk is not None + and int(chunk) > 0 + and int(chunk) < N + and not psicomp_chunked + ) + + if use_serial_chunks: + # CIP-0001 R14: serial map-reduce over row chunks for any kernel. + psi1_parts = [] + psi2 = None + psi0 = 0.0 + for sl in row_slices(N, chunk): + Xc = X[sl] + psi0 = psi0 + kern.psi0(Z, Xc).sum() + psi1_parts.append(kern.psi1(Z, Xc) * beta) + psi2_c = kern.psi2(Z, Xc) * beta + psi2 = psi2_c if psi2 is None else psi2 + psi2_c + psi1 = np.vstack(psi1_parts) + psi0 = psi0 * beta + elif uncertain_inputs: psi0 = kern.psi0(Z, X) psi1 = kern.psi1(Z, X)*beta psi2 = kern.psi2(Z, X)*beta + psi0 = psi0.sum()*beta else: psi0 = kern.Kdiag(X) psi1 = kern.K(X, Z) psi2 = np.dot(psi1.T,psi1)*beta psi1 = psi1*beta - - psi0 = psi0.sum()*beta + psi0 = psi0.sum()*beta if isinstance(Y, VariationalPosterior): m, s = Y.mean, Y.variance diff --git a/deepgp/layers/layers.py b/deepgp/layers/layers.py index de6951a..79e840f 100644 --- a/deepgp/layers/layers.py +++ b/deepgp/layers/layers.py @@ -21,9 +21,13 @@ def __init__(self, X, Y, Z, kernel, likelihood, normalizer=False, mpi_comm=None, mpi_root=0, - auto_update=True): + auto_update=True, + psi_chunk_size=None, + psi_threads=1): self.mpi_comm = mpi_comm self.mpi_root = mpi_root + self.psi_chunk_size = psi_chunk_size + self.psi_threads = 1 if psi_threads is None else int(psi_threads) self.psicov = False self.svi = False self.qU_ratio = 1. @@ -32,7 +36,7 @@ def __init__(self, X, Y, Z, kernel, likelihood, if inference_method is None: from ..inference import VarDTC_parallel, VarDTC if mpi_comm is None: - inference_method = VarDTC() + inference_method = VarDTC(psi_chunk_size=psi_chunk_size, psi_threads=self.psi_threads) else: inference_method = VarDTC_parallel(mpi_comm, mpi_root) elif inference_method=='inferentia' and mpi_comm is None: @@ -41,7 +45,7 @@ def __init__(self, X, Y, Z, kernel, likelihood, self.psicov = True elif inference_method=='svi': from ..inference import SVI_VarDTC - inference_method = SVI_VarDTC() + inference_method = SVI_VarDTC(psi_chunk_size=psi_chunk_size, psi_threads=self.psi_threads) self.svi = True super(SparseGP_MPI, self).__init__(X, Y, Z, kernel, @@ -50,6 +54,9 @@ def __init__(self, X, Y, Z, kernel, likelihood, inference_method=inference_method, name=name, Y_metadata=Y_metadata, normalizer=normalizer) + + from ..util.psi_chunking import apply_psi_chunking + apply_psi_chunking(self, psi_chunk_size, self.psi_threads) if self.svi: from ..util.misc import comp_mapping @@ -172,7 +179,8 @@ def __init__(self, layer_lower, Z=None, num_inducing=10, kernel=None, inference_method=None, uncertain_inputs=True, mpi_comm=None, mpi_root=0, back_constraint=True, - encoder=None, auto_update=True, name='layer'): + encoder=None, auto_update=True, name='layer', + psi_chunk_size=None, psi_threads=1): self.uncertain_inputs = uncertain_inputs self.layer_lower = layer_lower @@ -205,7 +213,9 @@ def __init__(self, layer_lower, super(Layer, self).__init__(X, Y, Z, kernel, likelihood, inference_method=inference_method, mpi_comm=mpi_comm, mpi_root=mpi_root, - auto_update=auto_update, name=name) + auto_update=auto_update, name=name, + psi_chunk_size=psi_chunk_size, + psi_threads=psi_threads) self.link_parameter(self.Kuu_sigma) if back_constraint: self.encoder = encoder @@ -304,7 +314,8 @@ def __init__(self, dim_down, dim_up, mpi_comm=None, mpi_root=0, back_constraint=True, encoder=None, auto_update=True, repeatX=False, - repeatXsplit=0, name='obslayer'): + repeatXsplit=0, name='obslayer', + psi_chunk_size=None, psi_threads=1): self.dim_up, self.dim_down = dim_up, dim_down self._Y = Y self.repeatX = repeatX @@ -321,7 +332,9 @@ def __init__(self, dim_down, dim_up, mpi_comm=mpi_comm, mpi_root=mpi_root, back_constraint=back_constraint, encoder=encoder, auto_update=auto_update, - name=name) + name=name, + psi_chunk_size=psi_chunk_size, + psi_threads=psi_threads) def set_as_toplayer(self, flag=True): if flag: @@ -339,7 +352,8 @@ def __init__(self, layer_lower, dim_up, kernel=None, inference_method=None, noise_var=1e-2, init='rand', mpi_comm=None, mpi_root=0, back_constraint=True, - encoder=None, auto_update=True, name='hiddenlayer'): + encoder=None, auto_update=True, name='hiddenlayer', + psi_chunk_size=None, psi_threads=1): self.dim_up, self.dim_down = dim_up, layer_lower.X.shape[1] #self.Y.shape[1] likelihood = likelihoods.Gaussian(variance=noise_var) @@ -354,7 +368,9 @@ def __init__(self, layer_lower, dim_up, mpi_comm=mpi_comm, mpi_root=mpi_root, back_constraint=back_constraint, encoder=encoder, auto_update=auto_update, - name=name) + name=name, + psi_chunk_size=psi_chunk_size, + psi_threads=psi_threads) def update_layer(self): super(HiddenLayer,self).update_layer() @@ -398,7 +414,8 @@ def __init__(self, layer_lower, dim_up, X=None, X_variance=None, Z=None, mpi_comm=None, mpi_root=0, encoder=None, back_constraint=True, - auto_update=True, name='tophiddenlayer'): + auto_update=True, name='tophiddenlayer', + psi_chunk_size=None, psi_threads=1): self.dim_up, self.dim_down = dim_up, layer_lower.X.shape[1] likelihood = likelihoods.Gaussian(variance=noise_var) @@ -411,7 +428,9 @@ def __init__(self, layer_lower, dim_up, X=None, X_variance=None, Z=None, mpi_comm=mpi_comm, mpi_root=mpi_root, back_constraint=back_constraint, encoder=encoder, auto_update=auto_update, - name=name) + name=name, + psi_chunk_size=psi_chunk_size, + psi_threads=psi_threads) def update_layer(self): super(TopHiddenLayer,self).update_layer() diff --git a/deepgp/models/model.py b/deepgp/models/model.py index 6a9743b..7d5b1c4 100644 --- a/deepgp/models/model.py +++ b/deepgp/models/model.py @@ -29,7 +29,8 @@ def __init__(self, nDims, Y, X=None, num_inducing=10, likelihood = None, inits='PCA', name='deepgp', kernels=None, obs_data='cont', back_constraint=True, encoder_dims=None, mpi_comm=None, mpi_root=0, repeatX=False, - inference_method=None, **kwargs): + inference_method=None, psi_chunk_size=None, psi_threads=1, + **kwargs): super(DeepGP, self).__init__(name=name) self._IN_OPTIMIZATION_ = False @@ -49,6 +50,8 @@ def __init__(self, nDims, Y, X=None, num_inducing=10, ) self.mpi_comm, self.mpi_root = mpi_comm, mpi_root + self.psi_chunk_size = psi_chunk_size + self.psi_threads = 1 if psi_threads is None else int(psi_threads) self.obs_data = obs_data self.nLayers = len(nDims)-1 self.nDims = nDims @@ -128,7 +131,9 @@ def __init__(self, nDims, Y, X=None, num_inducing=10, mpi_root=mpi_root, auto_update=auto_update, repeatX=True, - repeatXsplit=self.nDimsOrig[1])) + repeatXsplit=self.nDimsOrig[1], + psi_chunk_size=psi_chunk_size, + psi_threads=self.psi_threads)) self.layers[-1].X_dim_free = range(self.nDimsOrig[1]) self.layers[-1].X_dim_top = np.arange(self.nDimsOrig[1], self.nDimsOrig[1]+X.shape[1]).tolist() @@ -141,7 +146,9 @@ def __init__(self, nDims, Y, X=None, num_inducing=10, back_constraint=back_constraint, inference_method=inference_method, mpi_comm=mpi_comm, mpi_root=mpi_root, - auto_update=auto_update)) + auto_update=auto_update, + psi_chunk_size=psi_chunk_size, + psi_threads=self.psi_threads)) elif i==self.nLayers-1: if isinstance(X, VariationalPosterior): X_variance, uncertain_inputs = X.variance.values, True @@ -158,7 +165,9 @@ def __init__(self, nDims, Y, X=None, num_inducing=10, back_constraint=True if uncertain_inputs and back_constraint else False, inference_method=inference_method, mpi_comm=mpi_comm, mpi_root=mpi_root, - auto_update=auto_update)) + auto_update=auto_update, + psi_chunk_size=psi_chunk_size, + psi_threads=self.psi_threads)) else: self.layers.append(HiddenLayer(self.layers[-1], nDims[i+1], X=Xs[i], @@ -168,7 +177,9 @@ def __init__(self, nDims, Y, X=None, num_inducing=10, back_constraint=back_constraint, inference_method=inference_method, mpi_comm=mpi_comm, mpi_root=mpi_root, - auto_update=auto_update)) + auto_update=auto_update, + psi_chunk_size=psi_chunk_size, + psi_threads=self.psi_threads)) if self.nLayers==1: self.layers[0].set_as_toplayer() @@ -294,7 +305,9 @@ def add_layer(self, dim_up, X=None, mpi_root=self.mpi_root, back_constraint=back_constraint, X=X, X_variance=X_variance, Z=Z, - uncertain_inputs=uncertain_inputs)) + uncertain_inputs=uncertain_inputs, + psi_chunk_size=self.psi_chunk_size, + psi_threads=self.psi_threads)) ls.append(self.layers[-1]) self.nLayers += 1 self.nDims.append(dim_up) diff --git a/deepgp/testing/test_psi_chunking.py b/deepgp/testing/test_psi_chunking.py new file mode 100644 index 0000000..0dfa70b --- /dev/null +++ b/deepgp/testing/test_psi_chunking.py @@ -0,0 +1,74 @@ +# Copyright (c) 2015-2016, the authors (see AUTHORS.txt). +# Licensed under the BSD 3-clause license (see LICENSE.txt) + +"""CIP-0001 Phase 3: DeepGP psi_chunk_size / psi_threads wiring.""" + +import unittest + +import numpy as np +import GPy + +import deepgp + + +def _small_supervised(seed=0): + rng = np.random.RandomState(seed) + Y = rng.randn(24, 2) + X = rng.randn(24, 1) + kernels = [GPy.kern.RBF(3, ARD=True), GPy.kern.RBF(1, ARD=True)] + return [2, 3, 1], Y, X, kernels + + +class PsiChunkingTest(unittest.TestCase): + def test_chunked_matches_default_objective_and_grad(self): + nDims, Y, X, kernels = _small_supervised() + m0 = deepgp.DeepGP( + nDims, Y, X=X, kernels=kernels, num_inducing=4, back_constraint=True + ) + m1 = deepgp.DeepGP( + nDims, + Y, + X=X, + kernels=[k.copy() for k in kernels], + num_inducing=4, + back_constraint=True, + psi_chunk_size=7, + psi_threads=2, + ) + m1.param_array[:] = m0.param_array + m1.trigger_update() + self.assertAlmostEqual( + float(m0.log_likelihood()), float(m1.log_likelihood()), places=8 + ) + np.testing.assert_allclose(m0.gradient, m1.gradient, rtol=1e-8, atol=1e-8) + for layer in m1.layers: + self.assertEqual(layer.kern.psicomp.chunk_size, 7) + self.assertEqual(layer.kern.psicomp.n_threads, 2) + self.assertEqual(layer.inference_method.psi_chunk_size, 7) + + def test_checkgrad_with_chunking(self): + nDims, Y, X, kernels = _small_supervised(seed=1) + m = deepgp.DeepGP( + nDims, + Y, + X=X, + kernels=kernels, + num_inducing=4, + back_constraint=True, + psi_chunk_size=5, + psi_threads=2, + ) + self.assertTrue(m.checkgrad()) + + def test_defaults_leave_psicomp_unchunked(self): + nDims, Y, X, kernels = _small_supervised(seed=2) + m = deepgp.DeepGP( + nDims, Y, X=X, kernels=kernels, num_inducing=4, back_constraint=True + ) + for layer in m.layers: + self.assertIsNone(layer.kern.psicomp.chunk_size) + self.assertEqual(layer.kern.psicomp.n_threads, 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/deepgp/util/psi_chunking.py b/deepgp/util/psi_chunking.py new file mode 100644 index 0000000..ddebfd7 --- /dev/null +++ b/deepgp/util/psi_chunking.py @@ -0,0 +1,36 @@ +# Copyright (c) 2015-2016, the authors (see AUTHORS.txt). +# Licensed under the BSD 3-clause license (see LICENSE.txt) + +"""Helpers for CIP-0001 Phase 3 chunked psi-statistics.""" + +from __future__ import print_function + + +def row_slices(N, chunk_size): + """Yield contiguous row slices covering ``0..N-1``.""" + for start in range(0, N, int(chunk_size)): + yield slice(start, min(start + int(chunk_size), N)) + + +def configure_kern_psicomp(kern, chunk_size=None, n_threads=1): + """Apply chunking knobs to a kernel's ``psicomp`` if supported (GPy R12). + + Returns True when the kernel will honour ``chunk_size`` / ``n_threads``. + """ + psicomp = getattr(kern, "psicomp", None) + if psicomp is None or not hasattr(psicomp, "chunk_size"): + return False + psicomp.chunk_size = chunk_size + psicomp.n_threads = 1 if n_threads is None else int(n_threads) + return True + + +def apply_psi_chunking(layer, chunk_size=None, n_threads=1): + """Configure a layer's inference method and kernel for psi chunking.""" + inf = getattr(layer, "inference_method", None) + if inf is not None: + inf.psi_chunk_size = chunk_size + inf.psi_threads = 1 if n_threads is None else int(n_threads) + kern = getattr(layer, "kern", None) + if kern is not None: + configure_kern_psicomp(kern, chunk_size, n_threads)