make jax implementation modular and update timing info

Jaan Altosaar · Jaan Altosaar · commit f21223271a6a · 2021-05-25T17:31:33.000-04:00
diff --git a/.env b/.env
@@ -4,4 +4,6 @@
 JAX_PLATFORM_NAME=cpu
 
 # suppress tensorflow warnings
-TF_CPP_MIN_LOG_LEVEL=2
+TF_CPP_MIN_LOG_LEVEL=2
+
+TFDS_DATA_DIR=/scratch/gpfs/altosaar/tensorflow_datasets
diff --git a/README.md b/README.md
@@ -9,16 +9,20 @@ Variational inference is used to fit the model to binarized MNIST handwritten di
 
 Blog post: https://jaan.io/what-is-variational-autoencoder-vae-tutorial/
 
-Example output with importance sampling for estimating the marginal likelihood on Hugo Larochelle's Binary MNIST dataset. Final marginal likelihood on the test set was `-97.10` nats after 65k iterations.
+
+## PyTorch implementation
+
+(anaconda environment is in `environment-jax.yml`)
+
+Importance sampling is used to estimate the marginal likelihood on Hugo Larochelle's Binary MNIST dataset. The final marginal likelihood on the test set was `-97.10` nats is comparable to published numbers.
 
 ```
-$ python train_variational_autoencoder_pytorch.py --variational mean-field
-step:   0       train elbo: -558.28
-step:   0               valid elbo: -392.78     valid log p(x): -359.91
-step:   10000   train elbo: -106.67
-step:   10000           valid elbo: -109.12     valid log p(x): -103.11
-step:   20000   train elbo: -107.28
-step:   20000           valid elbo: -105.65     valid log p(x): -99.74
+$ python train_variational_autoencoder_pytorch.py --variational mean-field --use_gpu --data_dir $DAT --max_iterations 30000 --log_interval 10000
+Step 0          Train ELBO estimate: -558.027   Validation ELBO estimate: -384.432      Validation log p(x) estimate: -355.430  Speed: 2.72e+06 examples/s
+Step 10000      Train ELBO estimate: -111.323   Validation ELBO estimate: -109.048      Validation log p(x) estimate: -103.746  Speed: 2.64e+04 examples/s
+Step 20000      Train ELBO estimate: -103.013   Validation ELBO estimate: -107.655      Validation log p(x) estimate: -101.275  Speed: 2.63e+04 examples/s
+Step 29999      Test ELBO estimate: -106.642    Test log p(x) estimate: -100.309
+Total time: 2.49 minutes
 ```
 
 
@@ -36,12 +40,14 @@ step:   30000   train elbo: -98.70
 step:   30000           valid elbo: -103.76     valid log p(x): -97.71
 ```
 
-Using jax (anaconda environment is in `environment-jax.yml`):
+## jax implementation
+
+Using jax (anaconda environment is in `environment-jax.yml`), to get a 3x speedup over pytorch:
 ```
-Step 0          Train ELBO estimate: -565.785   Validation ELBO estimate: -565.775      Validation log p(x) estimate: -565.775  Speed: 3813003636 examples/s
-Step 10000      Train ELBO estimate: -99.048    Validation ELBO estimate: -105.412      Validation log p(x) estimate: -105.412  Speed: 134   examples/s
-Step 20000      Train ELBO estimate: -108.399   Validation ELBO estimate: -105.191      Validation log p(x) estimate: -105.191  Speed: 140   examples/s
-Step 30000      Train ELBO estimate: -100.839   Validation ELBO estimate: -105.404      Validation log p(x) estimate: -105.404  Speed: 139   examples/s
-Step 40000      Train ELBO estimate: -97.761    Validation ELBO estimate: -105.382      Validation log p(x) estimate: -105.382  Speed: 139   examples/s
-Step 50000      Train ELBO estimate: -98.228    Validation ELBO estimate: -105.718      Validation log p(x) estimate: -105.718  Speed: 139   examples/s
+$ python train_variational_autoencoder_jax.py --gpu 
+Step 0          Train ELBO estimate: -566.059   Validation ELBO estimate: -565.755      Validation log p(x) estimate: -557.914  Speed: 2.56e+11 examples/s
+Step 10000      Train ELBO estimate: -98.560    Validation ELBO estimate: -105.725      Validation log p(x) estimate: -98.973   Speed: 7.03e+04 examples/s
+Step 20000      Train ELBO estimate: -109.794   Validation ELBO estimate: -105.756      Validation log p(x) estimate: -97.914   Speed: 4.26e+04 examples/s
+Step 29999      Test ELBO estimate: -104.867    Test log p(x) estimate: -96.716
+Total time: 0.810 minutes
 ```
diff --git a/environment-pytorch.yml b/environment-pytorch.yml
@@ -0,0 +1,64 @@
+name: /scratch/gpfs/altosaar/environment-pytorch
+channels:
+  - pytorch
+  - nvidia
+  - defaults
+dependencies:
+  - _libgcc_mutex=0.1=main
+  - blas=1.0=mkl
+  - bzip2=1.0.8=h7b6447c_0
+  - ca-certificates=2021.4.13=h06a4308_1
+  - certifi=2020.12.5=py38h06a4308_0
+  - cudatoolkit=11.1.74=h6bb024c_0
+  - ffmpeg=4.3=hf484d3e_0
+  - freetype=2.10.4=h5ab3b9f_0
+  - gmp=6.2.1=h2531618_2
+  - gnutls=3.6.15=he1e5248_0
+  - h5py=2.10.0=py38hd6299e0_1
+  - hdf5=1.10.6=hb1b8bf9_0
+  - intel-openmp=2021.2.0=h06a4308_610
+  - jpeg=9b=h024ee3a_2
+  - lame=3.100=h7b6447c_0
+  - lcms2=2.12=h3be6417_0
+  - ld_impl_linux-64=2.33.1=h53a641e_7
+  - libffi=3.3=he6710b0_2
+  - libgcc-ng=9.1.0=hdf63c60_0
+  - libgfortran-ng=7.3.0=hdf63c60_0
+  - libiconv=1.15=h63c8f33_5
+  - libidn2=2.3.1=h27cfd23_0
+  - libpng=1.6.37=hbc83047_0
+  - libstdcxx-ng=9.1.0=hdf63c60_0
+  - libtasn1=4.16.0=h27cfd23_0
+  - libtiff=4.1.0=h2733197_1
+  - libunistring=0.9.10=h27cfd23_0
+  - libuv=1.40.0=h7b6447c_0
+  - lz4-c=1.9.3=h2531618_0
+  - mkl=2021.2.0=h06a4308_296
+  - mkl-service=2.3.0=py38h27cfd23_1
+  - mkl_fft=1.3.0=py38h42c9631_2
+  - mkl_random=1.2.1=py38ha9443f7_2
+  - ncurses=6.2=he6710b0_1
+  - nettle=3.7.2=hbbd107a_1
+  - ninja=1.10.2=hff7bd54_1
+  - numpy=1.20.2=py38h2d18471_0
+  - numpy-base=1.20.2=py38hfae3a4d_0
+  - olefile=0.46=py_0
+  - openh264=2.1.0=hd408876_0
+  - openssl=1.1.1k=h27cfd23_0
+  - pillow=8.2.0=py38he98fc37_0
+  - pip=21.1.1=py38h06a4308_0
+  - python=3.8.10=hdb3f193_7
+  - pytorch=1.8.1=py3.8_cuda11.1_cudnn8.0.5_0
+  - readline=8.1=h27cfd23_0
+  - setuptools=52.0.0=py38h06a4308_0
+  - six=1.15.0=py38h06a4308_0
+  - sqlite=3.35.4=hdfb4753_0
+  - tk=8.6.10=hbc83047_0
+  - torchaudio=0.8.1=py38
+  - torchvision=0.9.1=py38_cu111
+  - typing_extensions=3.7.4.3=pyha847dfd_0
+  - wheel=0.36.2=pyhd3eb1b0_0
+  - xz=5.2.5=h7b6447c_0
+  - zlib=1.2.11=h7b6447c_3
+  - zstd=1.4.9=haebb681_0
+prefix: /scratch/gpfs/altosaar/environment-pytorch
diff --git a/train_variational_autoencoder_jax.py b/train_variational_autoencoder_jax.py
@@ -8,11 +8,8 @@
 from calendar import c
 from typing import Generator, Mapping, NamedTuple, Sequence, Tuple
 
-import jax
 import numpy as np
-
-jax.config.update("jax_platform_name", "cpu")  # suppress warning about no GPUs
-
+import jax
 import haiku as hk
 import jax.numpy as jnp
 import optax
@@ -31,12 +28,11 @@ def add_args(parser):
     parser.add_argument("--hidden_size", type=int, default=512)
     parser.add_argument("--learning_rate", type=float, default=0.001)
     parser.add_argument("--batch_size", type=int, default=128)
-    parser.add_argument("--training_steps", type=int, default=100000)
+    parser.add_argument("--training_steps", type=int, default=30000)
     parser.add_argument("--log_interval", type=int, default=10000)
     parser.add_argument("--num_eval_samples", type=int, default=128)
     parser.add_argument("--gpu", default=False, action=argparse.BooleanOptionalAction)
     parser.add_argument("--random_seed", type=int, default=42)
-    parser.add_argument("--train_dir", type=pathlib.Path, default="/tmp")
 
 
 def load_dataset(
@@ -65,7 +61,7 @@ def __init__(
         hidden_size: int,
         output_shape: Sequence[int] = MNIST_IMAGE_SHAPE,
     ):
-        super().__init__()
+        super().__init__(name="model")
         self._latent_size = latent_size
         self._hidden_size = hidden_size
         self._output_shape = output_shape
@@ -93,7 +89,7 @@ class VariationalMeanField(hk.Module):
     """Mean field variational distribution q(z | x) parameterized by inference network."""
 
     def __init__(self, latent_size: int, hidden_size: int):
-        super().__init__()
+        super().__init__(name="variational")
         self._latent_size = latent_size
         self._hidden_size = hidden_size
         self.inference_network = hk.Sequential(
@@ -121,70 +117,49 @@ def __call__(self, x: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]:
         return q_z
 
 
-class ModelAndVariationalOutput(NamedTuple):
-    p_z: tfd.Distribution
-    p_x_given_z: tfd.Distribution
-    q_z: tfd.Distribution
-    z: jnp.ndarray
-
-
-class ModelAndVariational(hk.Module):
-    """Parent class for creating inputs to the variational inference algorithm."""
-
-    def __init__(self, latent_size: int, hidden_size: int, output_shape: Sequence[int]):
-        super().__init__()
-        self._latent_size = latent_size
-        self._hidden_size = hidden_size
-        self._output_shape = output_shape
-
-    def __call__(self, x: jnp.ndarray) -> ModelAndVariationalOutput:
-        x = x.astype(jnp.float32)
-        q_z = VariationalMeanField(self._latent_size, self._hidden_size)(x)
-        # use a single sample from variational distribution to train
-        # shape [num_samples, batch_size, latent_size]
-        z = q_z.sample(sample_shape=[1], seed=hk.next_rng_key())
-
-        p_z, p_x_given_z = Model(
-            self._latent_size, self._hidden_size, MNIST_IMAGE_SHAPE
-        )(x=x, z=z)
-        return ModelAndVariationalOutput(p_z, p_x_given_z, q_z, z)
-
-
 def main():
+    start_time = time.time()
     parser = argparse.ArgumentParser()
     add_args(parser)
     args = parser.parse_args()
-    model_and_variational = hk.transform(
-        lambda x: ModelAndVariational(
-            args.latent_size, args.hidden_size, MNIST_IMAGE_SHAPE
-        )(x)
+    rng_seq = hk.PRNGSequence(args.random_seed)
+    model = hk.transform(
+        lambda x, z: Model(args.latent_size, args.hidden_size, MNIST_IMAGE_SHAPE)(x, z)
+    )
+    variational = hk.transform(
+        lambda x: VariationalMeanField(args.latent_size, args.hidden_size)(x)
     )
+    p_params = model.init(
+        next(rng_seq),
+        np.zeros((1, *MNIST_IMAGE_SHAPE)),
+        np.zeros((1, args.latent_size)),
+    )
+    q_params = variational.init(next(rng_seq), np.zeros((1, *MNIST_IMAGE_SHAPE)))
+    params = hk.data_structures.merge(p_params, q_params)
+    optimizer = optax.rmsprop(args.learning_rate)
+    opt_state = optimizer.init(params)
 
-    # @jax.jit
+    @jax.jit
     def objective_fn(params: hk.Params, rng_key: PRNGKey, batch: Batch) -> jnp.ndarray:
         x = batch["image"]
-        out: ModelAndVariationalOutput = model_and_variational.apply(params, rng_key, x)
-        log_q_z = out.q_z.log_prob(out.z).sum(axis=-1)
+        predicate = lambda module_name, name, value: "model" in module_name
+        p_params, q_params = hk.data_structures.partition(predicate, params)
+        q_z = variational.apply(q_params, rng_key, x)
+        z = q_z.sample(sample_shape=[1], seed=rng_key)
+        p_z, p_x_given_z = model.apply(p_params, rng_key, x, z)
+        # out: ModelAndVariationalOutput = model_and_variational.apply(params, rng_key, x)
+        log_q_z = q_z.log_prob(z).sum(axis=-1)
         # sum over last three image dimensions (width, height, channels)
-        log_p_x_given_z = out.p_x_given_z.log_prob(x).sum(axis=(-3, -2, -1))
+        log_p_x_given_z = p_x_given_z.log_prob(x).sum(axis=(-3, -2, -1))
         # sum over latent dimension
-        log_p_z = out.p_z.log_prob(out.z).sum(axis=-1)
-
+        log_p_z = p_z.log_prob(z).sum(axis=-1)
         elbo = log_p_x_given_z + log_p_z - log_q_z
         # average elbo over number of samples
         elbo = elbo.mean(axis=0)
         # sum elbo over batch
         elbo = elbo.sum(axis=0)
         return -elbo
 
-    rng_seq = hk.PRNGSequence(args.random_seed)
-
-    params = model_and_variational.init(
-        next(rng_seq), np.zeros((1, *MNIST_IMAGE_SHAPE))
-    )
-    optimizer = optax.rmsprop(args.learning_rate)
-    opt_state = optimizer.init(params)
-
     @jax.jit
     def train_step(
         params: hk.Params, rng_key: PRNGKey, opt_state: optax.OptState, batch: Batch
@@ -201,13 +176,17 @@ def importance_weighted_estimate(
     ) -> Tuple[jnp.ndarray, jnp.ndarray]:
         """Estimate marginal log p(x) using importance sampling."""
         x = batch["image"]
-        out: ModelAndVariationalOutput = model_and_variational.apply(params, rng_key, x)
-        log_q_z = out.q_z.log_prob(out.z).sum(axis=-1)
+        # out: ModelAndVariationalOutput = model_and_variational.apply(params, rng_key, x)
+        predicate = lambda module_name, name, value: "model" in module_name
+        p_params, q_params = hk.data_structures.partition(predicate, params)
+        q_z = variational.apply(q_params, rng_key, x)
+        z = q_z.sample(args.num_eval_samples, seed=rng_key)
+        p_z, p_x_given_z = model.apply(p_params, rng_key, x, z)
+        log_q_z = q_z.log_prob(z).sum(axis=-1)
         # sum over last three image dimensions (width, height, channels)
-        log_p_x_given_z = out.p_x_given_z.log_prob(x).sum(axis=(-3, -2, -1))
+        log_p_x_given_z = p_x_given_z.log_prob(x).sum(axis=(-3, -2, -1))
         # sum over latent dimension
-        log_p_z = out.p_z.log_prob(out.z).sum(axis=-1)
-
+        log_p_z = p_z.log_prob(z).sum(axis=-1)
         elbo = log_p_x_given_z + log_p_z - log_q_z
         # importance sampling of approximate marginal likelihood with q(z)
         # as the proposal, and logsumexp in the sample dimension
@@ -253,15 +232,16 @@ def print_progress(step: int, examples_per_sec: float):
             f"Train ELBO estimate: {train_elbo:<5.3f}\t"
             f"Validation ELBO estimate: {elbo:<5.3f}\t"
             f"Validation log p(x) estimate: {log_p_x:<5.3f}\t"
-            f"Speed: {examples_per_sec:<5.0f} examples/s"
+            f"Speed: {examples_per_sec:<5.2e} examples/s"
         )
 
     t0 = time.time()
     for step in range(args.training_steps):
         if step % args.log_interval == 0:
-            examples_per_sec = args.log_interval / (time.time() - t0)
+            t1 = time.time()
+            examples_per_sec = args.log_interval * args.batch_size / (t1 - t0)
             print_progress(step, examples_per_sec)
-            t0 = time.time()
+            t0 = t1
         params, opt_state = train_step(params, next(rng_seq), opt_state, next(train_ds))
 
     test_ds = load_dataset(tfds.Split.TEST, args.batch_size, args.random_seed)
@@ -271,6 +251,7 @@ def print_progress(step: int, examples_per_sec: float):
         f"Test ELBO estimate: {elbo:<5.3f}\t"
         f"Test log p(x) estimate: {log_p_x:<5.3f}\t"
     )
+    print(f"Total time: {(time.time() - start_time) / 60:.3f} minutes")
 
 
 if __name__ == "__main__":
diff --git a/train_variational_autoencoder_pytorch.py b/train_variational_autoencoder_pytorch.py
