nn-bench.py

from time import perf_counter_ns
from scipy.stats import bootstrap
import numpy as np

def timer(f, *args, iter=1, confidence_level=0.95):
  runtimes_ns = np.zeros(iter)
  for i in range(iter):
    tic = perf_counter_ns()
    tmp = f(*args)
    toc = perf_counter_ns()
    runtimes_ns[i] = toc - tic
  runtimes_µs = runtimes_ns / 1000
  print(f"{f.__name__.ljust(25)}{np.mean(runtimes_µs):10.0f}µs", end="")
  if iter > 1:
    CI = bootstrap(
      (runtimes_µs,), np.mean, method="BCa", confidence_level=confidence_level
    ).confidence_interval
    print(f" (95% CI: [{CI.low:10.1f}, {CI.high:10.1f}])")
  else:
    print()
  return tmp

def dump(x, y):
  def shapeformat(shapestr):
    return shapestr.replace(",","][").replace("(","[").replace(")","]").replace(" ","")
  with open(shapeformat(f"./jax-{x.shape}f32-{y.shape}f32"), "wb") as f:
    futhark_data.dump(x, f, binary=True)
    futhark_data.dump(y, f, binary=True)

import jax
import jax.numpy as jnp
from jax import random
from jax.example_libraries import stax

import futhark_data


LATENT_DIM = 160
HIDDEN_DIM = 400
INPUT_DIM = 784

def main(num_iter):
  # with open("./mnist-f32", "rb") as f:
  #   x_np1 = next(futhark_data.load(f))
  #   x_np1 = x_np1.reshape(-1, INPUT_DIM)
  with open("./data/[60000][784]f32.in", "rb") as f:
    x_np1 = next(futhark_data.load(f))
    x_np1 = x_np1.reshape(-1, INPUT_DIM)

  # with open("./mnist-[60000][784]f32", "wb") as f:
  #   futhark_data.dump(x_np, f, binary=True)

  for batch_sz in [320]:
    x_np = x_np1[:batch_sz]
    print()
    print("Batch size", batch_sz)

    # with open(f"./jax-mnist-[{batch_sz}][784]f32", "wb") as f:
    #   futhark_data.dump(x_np, f, binary=True)

    # measure JAX device transfer time
    timer(jax.device_put, x_np)
    x = jax.device_put(x_np)

    (init, fwd) = stax.serial(
      stax.Dense(HIDDEN_DIM, W_init=stax.randn()),
      stax.Softplus,
    )
    (_input_shape, params) = init(random.PRNGKey(0), x.shape)

    matmul_weights = params[0][0]
    dump(x_np, matmul_weights)
    dump(x_np, matmul_weights.T)

    # Time matrix multiplication.
    def matmul(x, W):
      return jnp.dot(x, W)

    matmul_jit_ = jax.jit(matmul)
    def matmul_jit(x, W):
      return matmul_jit_(x, W).block_until_ready()
    timer(matmul_jit, x, matmul_weights)
    timer(matmul_jit, x, matmul_weights, iter=num_iter)

    # Time weight initialisation, because this is tedious to isolate in Futhark.
    def W_init(x):
      (init, fwd) = stax.serial(
        stax.Dense(HIDDEN_DIM, W_init=stax.randn()),
        stax.Softplus,
      )
      (input_shape, params) = init(random.PRNGKey(0), x.shape)
      return params

    W_init_jit_ = jax.jit(W_init)
    def W_init_jit(x):
      return W_init_jit_(x)
    timer(W_init_jit, x, iter=1)
    timer(W_init_jit, x, iter=num_iter)

    # Time one_layer neural network including weight initialisation.
    def one_layer(x):
      (init, fwd) = stax.serial(
        stax.Dense(HIDDEN_DIM, W_init=stax.randn()),
        stax.Softplus,
      )
      (_input_shape, params) = init(random.PRNGKey(0), x.shape)
      return fwd(params, x)

    one_layer_jit_ = jax.jit(one_layer)
    def one_layer_jit(x):
      return one_layer_jit_(x).block_until_ready()
    timer(one_layer_jit, x, iter=1)
    timer(one_layer_jit, x, iter=num_iter)

    # AD
    def vjp_fun(x, dy):
      _, f_vjp = jax.vjp(one_layer, x)
      xbar, = f_vjp(dy)
      return xbar

    vjp_one_layer_jit_ = jax.jit(vjp_fun)
    def vjp_one_layer_jit(x, dy):
      return vjp_one_layer_jit_(x, dy).block_until_ready()
    dy = jax.random.normal(jax.random.PRNGKey(1337), (x.shape[0], HIDDEN_DIM))
    timer(vjp_one_layer_jit, x, dy, iter=1)
    timer(vjp_one_layer_jit, x, dy, iter=num_iter)

    # Encoder
    def f(x):
      (init, fwd) = stax.serial(
          stax.Dense(HIDDEN_DIM, W_init=stax.randn()),
          stax.Softplus,
          stax.FanOut(2),
          stax.parallel(
              stax.Dense(LATENT_DIM, W_init=stax.randn()),
              stax.serial(stax.Dense(LATENT_DIM, W_init=stax.randn()), stax.Exp),
          ),
      )
      (_input_shape, params) = init(random.PRNGKey(0), x.shape)
      return fwd(params, x)

    f_jit = jax.jit(f)
    def encoder_jit(x):
      (a,b) = f_jit(x)
      return (a.block_until_ready(), b.block_until_ready())
    timer(encoder_jit, x, iter=1)
    timer(encoder_jit, x, iter=num_iter)

    def vjp_f(x, dy):
      _, f_vjp = jax.vjp(f, x)
      xbar, = f_vjp(dy)
      return xbar

    vjp_f_jit = jax.jit(vjp_f)
    def vjp_encoder_jit(x, dy):
      return vjp_f_jit(x, dy).block_until_ready()
    dy = [jax.random.normal(jax.random.PRNGKey(1337), (x.shape[0], LATENT_DIM)),
          jax.random.normal(jax.random.PRNGKey(42), (x.shape[0], LATENT_DIM))]
    timer(vjp_encoder_jit, x, dy, iter=1)
    timer(vjp_encoder_jit, x, dy, iter=num_iter)

    # Decoder
    def f(x):
      (init, fwd) = stax.serial(
        stax.Dense(HIDDEN_DIM, W_init=stax.randn()),
        stax.Softplus,
        stax.Dense(INPUT_DIM, W_init=stax.randn()),
        stax.Sigmoid,
      )
      (_input_shape, params) = init(random.PRNGKey(0), x.shape)
      return fwd(params, x)

    f_jit = jax.jit(f)
    def decoder_jit(x):
      return f_jit(x).block_until_ready()
    timer(decoder_jit, x, iter=1)
    timer(decoder_jit, x, iter=num_iter)

    def vjp_f(x, dy):
      _, f_vjp = jax.vjp(f, x)
      xbar, = f_vjp(dy)
      return xbar

    vjp_f_jit = jax.jit(vjp_f)
    def vjp_decoder_jit(x, dy):
      return vjp_f_jit(x, dy).block_until_ready()
    dy = jax.random.normal(jax.random.PRNGKey(1337), (x.shape[0], INPUT_DIM))
    timer(vjp_decoder_jit, x, dy, iter=1)
    timer(vjp_decoder_jit, x, dy, iter=num_iter)

if __name__ == "__main__":
  main(300)