Benchmarking NumPy and TensorFlow2 on simple network model

np_tf_bench.py

import time
import tqdm
import numpy as np
import tensorflow as tf
import numba as nb
#tf.config.experimental.set_visible_devices([], 'GPU')
tf.zeros((1,), tf.float32)



def bench(nr, inner_size):
    nt, nn = 10000, 96
    a, tau, k, dt = 0.5, 3.0, 0.1, 0.01

    def run_np_inner(x, y):
        for t in range(inner_size):
            mf = x.mean(axis=0)
            dx = tau * (x - x**3/3 + y)
            dy = (1/tau) * (a - x + mf)
            x = x + dt * dx
            y = y + dt * dy
        return x, y

    def run_np(x, y):
        for t in range(nt//inner_size):
            x, y = run_np_inner(x, y)
        return x, y

    @tf.function
    def run_tf_inner(x, y):
        for t in range(inner_size):
            mf = tf.reduce_mean(x, axis=0)
            dx = tau * (x - x**3/3 + y)
            dy = (1/tau) * (a - x + mf)
            x = x + dt * dx
            y = y + dt * dy
        return x, y

    def run_tf(x, y):
        for t in range(nt//inner_size):
            x, y = run_tf_inner(x, y)
        return x, y

    # numpy
    x, y = np.zeros((2, nn, nr), 'f')
    tic = time.time()
    run_np(x, y)
    toc_np = time.time() - tic

    # tf
    x, y = tf.zeros((2, nn, nr), tf.float32)
    tic = time.time()
    run_tf(x, y)
    toc_tf = time.time() - tic

    return toc_tf / toc_np

nr = 1024
ni = 20
print("%d %d %0.3f" % (ni, nr, bench(nr, ni)))

# 10885H/RTX4000 on "better battery"
# GPU nr=1024 ni=20 GPU 0.07 CPU 0.18
# GPU nr=64 ni=20  GPU 0.81 CPU 0.56
# GPU nr=32 ni=20  GPU 1.12 CPU 0.87
# GPU nr= 8 ni=20  GPU 2.38 CPU 1.53
# GPU nr= 1 ni=20  GPU 4.32 CPU 2.57
# GPU nr= 1 ni=10  GPU      CPU 1.77
# GPU nr= 1 ni=13  GPU 3.99 CPU