ppl/nn_comparison.py at main · nhey/ppl · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
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__}:\t{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


HIDDEN_DIM = 400
INPUT_DIM = 784

def main():
  with open("./mnist-f32", "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, 3200, 32000, 60000]:
    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)

    # fwd_jit_ = jax.jit(fwd)
    # def fwd_jit(params, x):
    #   return fwd_jit_(params, x).block_until_ready()
    # # measure JAX compilation time
    # timer(fwd_jit, params, x)
    # # measure JAX runtime
    # timer(fwd_jit, params, x, iter=10)

    # Time weight initialisation, because this is tedious to isolate in Futhark.
    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=10)

    # Time entire program (more similar to Futhark).
    def program(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)

    program_jit_ = jax.jit(program)
    def program_jit(x):
      return program_jit_(x).block_until_ready()

    timer(program_jit, x, iter=1)
    timer(program_jit, x, iter=10)

if __name__ == "__main__":
  main()