Hamiltonian Monte Carlo with Energy Conserving Subsampling

docs/source/conf.py

@@ -9,7 +9,6 @@
 import nbsphinx
 import sphinx_rtd_theme
 
-
 # import pkg_resources
 
 # -*- coding: utf-8 -*-
@@ -33,6 +32,7 @@
 
 # HACK: This is to ensure that local functions are documented by sphinx.
 from numpyro.infer.hmc import hmc  # noqa: E402
+
 hmc(None, None)
 
 # -- Project information -----------------------------------------------------

examples/covtype.py

@@ -4,13 +4,18 @@
 import argparse
 import time
 
+import matplotlib.pyplot as plt
+
 from jax import random
 import jax.numpy as jnp
 
 import numpyro
 import numpyro.distributions as dist
 from numpyro.examples.datasets import COVTYPE, load_dataset
-from numpyro.infer import MCMC, NUTS
+from numpyro.infer import HMC, HMCECS, MCMC, NUTS, SA, SVI, Trace_ELBO, init_to_value
+from numpyro.infer.autoguide import AutoBNAFNormal
+from numpyro.infer.hmc_gibbs import taylor_proxy
+from numpyro.infer.reparam import NeuTraReparam
 
 
 def _load_dataset():
@@ -33,22 +38,76 @@ def _load_dataset():
     return features, labels
 
 
-def model(data, labels):
+def model(data, labels, subsample_size=None):
     dim = data.shape[1]
     coefs = numpyro.sample('coefs', dist.Normal(jnp.zeros(dim), jnp.ones(dim)))
-    logits = jnp.dot(data, coefs)
-    return numpyro.sample('obs', dist.Bernoulli(logits=logits), obs=labels)
+    with numpyro.plate("N", data.shape[0], subsample_size=subsample_size) as idx:
+        logits = jnp.dot(data[idx], coefs)
+        return numpyro.sample('obs', dist.Bernoulli(logits=logits), obs=labels[idx])
 
 
 def benchmark_hmc(args, features, labels):
-    step_size = jnp.sqrt(0.5 / features.shape[0])
-    trajectory_length = step_size * args.num_steps
     rng_key = random.PRNGKey(1)
     start = time.time()
-    kernel = NUTS(model, trajectory_length=trajectory_length)
-    mcmc = MCMC(kernel, 0, args.num_samples)
-    mcmc.run(rng_key, features, labels)
-    mcmc.print_summary()
+    # a MAP estimate at the following source
+    # https://github.com/google/edward2/blob/master/examples/no_u_turn_sampler/logistic_regression.py#L117
+    ref_params = {"coefs": jnp.array([
+        +2.03420663e+00, -3.53567265e-02, -1.49223924e-01, -3.07049364e-01,
+        -1.00028366e-01, -1.46827862e-01, -1.64167881e-01, -4.20344204e-01,
+        +9.47479829e-02, -1.12681836e-02, +2.64442056e-01, -1.22087866e-01,
+        -6.00568838e-02, -3.79419506e-01, -1.06668741e-01, -2.97053963e-01,
+        -2.05253899e-01, -4.69537191e-02, -2.78072730e-02, -1.43250525e-01,
+        -6.77954629e-02, -4.34899796e-03, +5.90927452e-02, +7.23133609e-02,
+        +1.38526391e-02, -1.24497898e-01, -1.50733739e-02, -2.68872194e-02,
+        -1.80925727e-02, +3.47936489e-02, +4.03552800e-02, -9.98773426e-03,
+        +6.20188080e-02, +1.15002751e-01, +1.32145107e-01, +2.69109547e-01,
+        +2.45785132e-01, +1.19035013e-01, -2.59744357e-02, +9.94279515e-04,
+        +3.39266285e-02, -1.44057125e-02, -6.95222765e-02, -7.52013028e-02,
+        +1.21171586e-01, +2.29205526e-02, +1.47308692e-01, -8.34354162e-02,
+        -9.34122875e-02, -2.97472421e-02, -3.03937674e-01, -1.70958012e-01,
+        -1.59496680e-01, -1.88516974e-01, -1.20889175e+00])}
+    if args.algo == "HMC":
+        step_size = jnp.sqrt(0.5 / features.shape[0])
+        trajectory_length = step_size * args.num_steps
+        kernel = HMC(model, step_size=step_size, trajectory_length=trajectory_length, adapt_step_size=False,
+                     dense_mass=args.dense_mass)
+        subsample_size = None
+    elif args.algo == "NUTS":
+        kernel = NUTS(model, dense_mass=args.dense_mass)
+        subsample_size = None
+    elif args.algo == "HMCECS":
+        subsample_size = 1000
+        inner_kernel = NUTS(model, init_strategy=init_to_value(values=ref_params),
+                            dense_mass=args.dense_mass)
+        # note: if num_blocks=100, we'll update 10 index at each MCMC step
+        # so it took 50000 MCMC steps to iterative the whole dataset
+        kernel = HMCECS(inner_kernel, num_blocks=100, proxy=taylor_proxy(ref_params))
+    elif args.algo == "SA":
+        # NB: this kernel requires large num_warmup and num_samples
+        # and running on GPU is much faster than on CPU
+        kernel = SA(model, adapt_state_size=1000, init_strategy=init_to_value(values=ref_params))
+        subsample_size = None
+    elif args.algo == "FlowHMCECS":
+        subsample_size = 1000
+        guide = AutoBNAFNormal(model, num_flows=1, hidden_factors=[8])
+        svi = SVI(model, guide, numpyro.optim.Adam(0.01), Trace_ELBO())
+        params, losses = svi.run(random.PRNGKey(2), 2000, features, labels)
+        plt.plot(losses)
+        plt.show()
+
+        neutra = NeuTraReparam(guide, params)
+        neutra_model = neutra.reparam(model)
+        neutra_ref_params = {"auto_shared_latent": jnp.zeros(55)}
+        # no need to adapt mass matrix if the flow does a good job
+        inner_kernel = NUTS(neutra_model, init_strategy=init_to_value(values=neutra_ref_params),
+                            adapt_mass_matrix=False)
+        kernel = HMCECS(inner_kernel, num_blocks=100, proxy=taylor_proxy(neutra_ref_params))
+    else:
+        raise ValueError("Invalid algorithm, either 'HMC', 'NUTS', or 'HMCECS'.")
+    mcmc = MCMC(kernel, args.num_warmup, args.num_samples)
+    mcmc.run(rng_key, features, labels, subsample_size, extra_fields=("accept_prob",))
+    print("Mean accept prob:", jnp.mean(mcmc.get_extra_fields()["accept_prob"]))
+    mcmc.print_summary(exclude_deterministic=False)
     print('\nMCMC elapsed time:', time.time() - start)
 
 
@@ -60,14 +119,20 @@ def main(args):
 if __name__ == '__main__':
     assert numpyro.__version__.startswith('0.5.0')
     parser = argparse.ArgumentParser(description="parse args")
-    parser.add_argument('-n', '--num-samples', default=100, type=int, help='number of samples')
+    parser.add_argument('-n', '--num-samples', default=1000, type=int, help='number of samples')
+    parser.add_argument('--num-warmup', default=1000, type=int, help='number of warmup steps')
     parser.add_argument('--num-steps', default=10, type=int, help='number of steps (for "HMC")')
     parser.add_argument('--num-chains', nargs='?', default=1, type=int)
-    parser.add_argument('--algo', default='NUTS', type=str, help='whether to run "HMC" or "NUTS"')
-    parser.add_argument('--device', default='cpu', type=str, help='use "cpu" or "gpu".')
+    parser.add_argument('--algo', default='HMCECS', type=str,
+                        help='whether to run "HMCECS", "NUTS", "HMCECS", "SA" or "FlowHMCECS"')
+    parser.add_argument('--dense-mass', action="store_true")
+    parser.add_argument('--x64', action="store_true")
+    parser.add_argument('--device', default='gpu', type=str, help='use "cpu" or "gpu".')
     args = parser.parse_args()
 
     numpyro.set_platform(args.device)
     numpyro.set_host_device_count(args.num_chains)
+    if args.x64:
+        numpyro.enable_x64()
 
     main(args)

numpyro/diagnostics.py

@@ -161,6 +161,7 @@ def effective_sample_size(x):
     :return: effective sample size of ``x``.
     :rtype: numpy.ndarray
     """
+
     assert x.ndim >= 2
     assert x.shape[1] >= 2
 

numpyro/distributions/continuous.py

@@ -962,7 +962,7 @@ def sample(self, key, sample_shape=()):
 
     @validate_sample
     def log_prob(self, value):
-        normalize_term = jnp.log(jnp.sqrt(2 * jnp.pi) * self.scale)
+        normalize_term = jnp.log(jnp.sqrt(2 * jnp.pi) * self.scale)  # TODO:Added jnp.abs
         value_scaled = (value - self.loc) / self.scale
         return -0.5 * value_scaled ** 2 - normalize_term
 

numpyro/examples/datasets.py

@@ -4,11 +4,13 @@
 from collections import namedtuple
 import csv
 import gzip
+import io
 import os
 import pickle
 import struct
 from urllib.parse import urlparse
 from urllib.request import urlretrieve
+import warnings
 import zipfile
 
 import numpy as np
@@ -23,52 +25,47 @@
                                             '.data'))
 os.makedirs(DATA_DIR, exist_ok=True)
 
-
 dset = namedtuple('dset', ['name', 'urls'])
 
-
 BASEBALL = dset('baseball', [
     'https://d2hg8soec8ck9v.cloudfront.net/datasets/EfronMorrisBB.txt',
 ])
 
-
 COVTYPE = dset('covtype', [
     'https://d2hg8soec8ck9v.cloudfront.net/datasets/covtype.zip',
 ])
 
-
 DIPPER_VOLE = dset('dipper_vole', [
     'https://github.com/pyro-ppl/datasets/blob/master/dipper_vole.zip?raw=true',
 ])
 
-
 MNIST = dset('mnist', [
     'https://d2hg8soec8ck9v.cloudfront.net/datasets/mnist/train-images-idx3-ubyte.gz',
     'https://d2hg8soec8ck9v.cloudfront.net/datasets/mnist/train-labels-idx1-ubyte.gz',
     'https://d2hg8soec8ck9v.cloudfront.net/datasets/mnist/t10k-images-idx3-ubyte.gz',
     'https://d2hg8soec8ck9v.cloudfront.net/datasets/mnist/t10k-labels-idx1-ubyte.gz',
 ])
 
-
 SP500 = dset('SP500', [
     'https://d2hg8soec8ck9v.cloudfront.net/datasets/SP500.csv',
 ])
 
-
 UCBADMIT = dset('ucbadmit', [
     'https://d2hg8soec8ck9v.cloudfront.net/datasets/UCBadmit.csv',
 ])
 
-
 LYNXHARE = dset('lynxhare', [
     'https://d2hg8soec8ck9v.cloudfront.net/datasets/LynxHare.txt',
 ])
 
-
 JSB_CHORALES = dset('jsb_chorales', [
     'https://d2hg8soec8ck9v.cloudfront.net/datasets/polyphonic/jsb_chorales.pickle',
 ])
 
+HIGGS = dset("higgs", [
+    "https://archive.ics.uci.edu/ml/machine-learning-databases/00280/HIGGS.csv.gz",
+])
+
 
 def _download(dset):
     for url in dset.urls:
@@ -86,7 +83,7 @@ def _load_baseball():
     def train_test_split(file):
         train, test, player_names = [], [], []
         with open(file, 'r') as f:
-            csv_reader = csv.DictReader(f, delimiter='\t', quoting=csv.QUOTE_NONE)
+            csv_reader = csv.reader(f, delimiter='\t', quoting=csv.QUOTE_NONE)
             for row in csv_reader:
                 player_names.append(row['FirstName'] + ' ' + row['LastName'])
                 at_bats, hits = row['At-Bats'], row['Hits']
@@ -240,6 +237,21 @@ def _load_jsb_chorales():
     return processed_dataset
 
 
+def _load_higgs():
+    warnings.warn("Higgs is a 2.6 GB dataset")
+    _download(HIGGS)
+
+    file_path = os.path.join(DATA_DIR, 'HIGGS.csv.gz')
+    with io.TextIOWrapper(gzip.open(file_path, 'rb')) as f:
+        csv_reader = csv.reader(f, delimiter=',', quoting=csv.QUOTE_NONE)
+        obs = []
+        data = []
+        for row in csv_reader:
+            obs.append(row[0])
+            data.append(row[1:])
+    return np.stack(obs), np.stack(data)
+
+
 def _load(dset):
     if dset == BASEBALL:
         return _load_baseball()
@@ -257,6 +269,8 @@ def _load(dset):
         return _load_lynxhare()
     elif dset == JSB_CHORALES:
         return _load_jsb_chorales()
+    elif dset == HIGGS:
+        return _load_higgs()
     raise ValueError('Dataset - {} not found.'.format(dset.name))
 
 

numpyro/handlers.py

@@ -1,6 +1,5 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
-
 """
 This provides a small set of effect handlers in NumPyro that are modeled
 after Pyro's `poutine <http://docs.pyro.ai/en/stable/poutine.html>`_ module.
@@ -136,6 +135,7 @@ class trace(Messenger):
                       'type': 'sample',
                       'value': DeviceArray(-0.20584235, dtype=float32)})])
     """
+
     def __enter__(self):
         super(trace, self).__enter__()
         self.trace = OrderedDict()
@@ -146,7 +146,7 @@ def postprocess_message(self, msg):
             # skip recording helper messages e.g. `control_flow`, `to_data`, `to_funsor`
             # which has no name
             return
-        assert not(msg['type'] == 'sample' and msg['name'] in self.trace), \
+        assert not (msg['type'] == 'sample' and msg['name'] in self.trace), \
             'all sites must have unique names but got `{}` duplicated'.format(msg['name'])
         self.trace[msg['name']] = msg.copy()
 
@@ -191,6 +191,7 @@ class replay(Messenger):
        -0.20584235
        >>> assert replayed_trace['a']['value'] == exec_trace['a']['value']
     """
+
     def __init__(self, fn=None, guide_trace=None):
         assert guide_trace is not None
         self.guide_trace = guide_trace
@@ -234,6 +235,7 @@ class block(Messenger):
        >>> assert 'a' not in trace_block_a
        >>> assert 'b' in trace_block_a
     """
+
     def __init__(self, fn=None, hide_fn=None, hide=None):
         if hide_fn is not None:
             self.hide_fn = hide_fn
@@ -350,6 +352,7 @@ class condition(Messenger):
        >>> assert exec_trace['a']['value'] == -1
        >>> assert exec_trace['a']['is_observed']
     """
+
     def __init__(self, fn=None, data=None, condition_fn=None):
         self.condition_fn = condition_fn
         self.data = data
@@ -386,6 +389,7 @@ class infer_config(Messenger):
     :param fn: a stochastic function (callable containing NumPyro primitive calls)
     :param config_fn: a callable taking a site and returning an infer dict
     """
+
     def __init__(self, fn=None, config_fn=None):
         super().__init__(fn)
         self.config_fn = config_fn
@@ -470,6 +474,7 @@ class mask(Messenger):
     :param mask: a boolean or a boolean-valued array for masking elementwise log
         probability of sample sites (`True` includes a site, `False` excludes a site).
     """
+
     def __init__(self, fn=None, mask=True):
         if lax.dtype(mask) != 'bool':
             raise ValueError("`mask` should be a bool array.")
@@ -506,6 +511,7 @@ class reparam(Messenger):
         :class:`~numpyro.infer.reparam.Reparam` or None.
     :type config: dict or callable
     """
+
     def __init__(self, fn=None, config=None):
         assert isinstance(config, dict) or callable(config)
         self.config = config
@@ -550,6 +556,7 @@ class scale(Messenger):
         of log probability.
     :type scale: float or numpy.ndarray
     """
+
     def __init__(self, fn=None, scale=1.):
         if not_jax_tracer(scale):
             if np.any(np.less_equal(scale, 0)):
@@ -587,6 +594,7 @@ class scope(Messenger):
     :param str prefix: a string to prepend to sample names
     :param str divider: a string to join the prefix and sample name; default to `'/'`
     """
+
     def __init__(self, fn=None, prefix='', divider='/'):
         self.prefix = prefix
         self.divider = divider
@@ -638,6 +646,7 @@ class seed(Messenger):
        >>> y = handlers.seed(model, rng_seed=1)()
        >>> assert x == y
     """
+
     def __init__(self, fn=None, rng_seed=None):
         if isinstance(rng_seed, int) or (isinstance(rng_seed, jnp.ndarray) and not jnp.shape(rng_seed)):
             rng_seed = random.PRNGKey(rng_seed)
@@ -647,10 +656,10 @@ def __init__(self, fn=None, rng_seed=None):
         super(seed, self).__init__(fn)
 
     def process_message(self, msg):
-        if (msg['type'] == 'sample' and not msg['is_observed'] and
-                msg['kwargs']['rng_key'] is None) or msg['type'] in ['prng_key', 'plate', 'control_flow']:
-            # no need to create a new key when value is available
+        if (msg['type'] == 'sample' and not msg['is_observed'] and msg['kwargs']['rng_key'] is None) \
+                or msg['type'] in ['prng_key', 'plate', 'control_flow']:
             if msg['value'] is not None:
+                # no need to create a new key when value is available
                 return
             self.rng_key, rng_key_sample = random.split(self.rng_key)
             msg['kwargs']['rng_key'] = rng_key_sample
@@ -691,6 +700,7 @@ class substitute(Messenger):
        >>> exec_trace = trace(substitute(model, {'a': -1})).get_trace()
        >>> assert exec_trace['a']['value'] == -1
     """
+
     def __init__(self, fn=None, data=None, substitute_fn=None):
         self.substitute_fn = substitute_fn
         self.data = data
@@ -760,6 +770,7 @@ class do(Messenger):
       >>> assert not exec_trace['z'].get('stop', None)
       >>> assert z_square == 1
     """
+
     def __init__(self, fn=None, data=None):
         self.data = data
         self._intervener_id = str(id(self))