Source code for nengo_dl.benchmarks

"""
Benchmark networks and utilities for evaluating NengoDL's performance.
"""

import inspect
import itertools
import os
import random
import time

import click
import matplotlib.pyplot as plt
import nengo
import numpy as np
import tensorflow as tf

import nengo_dl
from nengo_dl.compat import tf_compat


[docs]def cconv(dimensions, neurons_per_d, neuron_type): """ Circular convolution (EnsembleArray) benchmark. Parameters ---------- dimensions : int Number of dimensions for vector values neurons_per_d : int Number of neurons to use per vector dimension neuron_type : `~nengo.neurons.NeuronType` Simulation neuron type Returns ------- net : `nengo.Network` benchmark network """ with nengo.Network(label="cconv", seed=0) as net: net.config[nengo.Ensemble].neuron_type = neuron_type net.config[nengo.Ensemble].gain = nengo.dists.Choice([1, -1]) net.config[nengo.Ensemble].bias = nengo.dists.Uniform(-1, 1) net.cconv = nengo.networks.CircularConvolution(neurons_per_d, dimensions) net.inp_a = nengo.Node([0] * dimensions) net.inp_b = nengo.Node([1] * dimensions) nengo.Connection(net.inp_a, net.cconv.A) nengo.Connection(net.inp_b, net.cconv.B) net.p = nengo.Probe(net.cconv.output) return net
[docs]def integrator(dimensions, neurons_per_d, neuron_type): """ Single integrator ensemble benchmark. Parameters ---------- dimensions : int Number of dimensions for vector values neurons_per_d : int Number of neurons to use per vector dimension neuron_type : `~nengo.neurons.NeuronType` Simulation neuron type Returns ------- net : `nengo.Network` benchmark network """ with nengo.Network(label="integrator", seed=0) as net: net.config[nengo.Ensemble].neuron_type = neuron_type net.config[nengo.Ensemble].gain = nengo.dists.Choice([1, -1]) net.config[nengo.Ensemble].bias = nengo.dists.Uniform(-1, 1) net.integ = nengo.networks.EnsembleArray(neurons_per_d, dimensions) nengo.Connection(net.integ.output, net.integ.input, synapse=0.01) net.inp = nengo.Node([0] * dimensions) nengo.Connection(net.inp, net.integ.input, transform=0.01) net.p = nengo.Probe(net.integ.output) return net
[docs]def pes(dimensions, neurons_per_d, neuron_type): """ PES learning rule benchmark. Parameters ---------- dimensions : int Number of dimensions for vector values neurons_per_d : int Number of neurons to use per vector dimension neuron_type : `~nengo.neurons.NeuronType` Simulation neuron type Returns ------- net : `nengo.Network` benchmark network """ with nengo.Network(label="pes", seed=0) as net: net.config[nengo.Ensemble].neuron_type = neuron_type net.config[nengo.Ensemble].gain = nengo.dists.Choice([1, -1]) net.config[nengo.Ensemble].bias = nengo.dists.Uniform(-1, 1) net.inp = nengo.Node([1] * dimensions) net.pre = nengo.Ensemble(neurons_per_d * dimensions, dimensions) net.post = nengo.Node(size_in=dimensions) nengo.Connection(net.inp, net.pre) conn = nengo.Connection(net.pre, net.post, learning_rule_type=nengo.PES()) nengo.Connection(net.post, conn.learning_rule, transform=-1) nengo.Connection(net.inp, conn.learning_rule) net.p = nengo.Probe(net.post) return net
[docs]def basal_ganglia(dimensions, neurons_per_d, neuron_type): """ Basal ganglia network benchmark. Parameters ---------- dimensions : int Number of dimensions for vector values neurons_per_d : int Number of neurons to use per vector dimension neuron_type : `~nengo.neurons.NeuronType` Simulation neuron type Returns ------- net : `nengo.Network` benchmark network """ with nengo.Network(label="basal_ganglia", seed=0) as net: net.config[nengo.Ensemble].neuron_type = neuron_type net.inp = nengo.Node([1] * dimensions) net.bg = nengo.networks.BasalGanglia(dimensions, neurons_per_d) nengo.Connection(net.inp, net.bg.input) net.p = nengo.Probe(net.bg.output) return net
[docs]def mnist(use_tensor_layer=True): """ A network designed to stress-test tensor layers (based on mnist net). Parameters ---------- use_tensor_layer : bool If True, use individual tensor_layers to build the network, as opposed to a single TensorNode containing all layers. Returns ------- net : `nengo.Network` benchmark network """ with nengo.Network() as net: # create node to feed in images net.inp = nengo.Node(np.ones(28 * 28)) if use_tensor_layer: nengo_nl = nengo.RectifiedLinear() ensemble_params = dict( max_rates=nengo.dists.Choice([100]), intercepts=nengo.dists.Choice([0]) ) amplitude = 1 synapse = None x = nengo_dl.tensor_layer( net.inp, tf_compat.layers.conv2d, shape_in=(28, 28, 1), filters=32, kernel_size=3, ) x = nengo_dl.tensor_layer(x, nengo_nl, **ensemble_params) x = nengo_dl.tensor_layer( x, tf_compat.layers.conv2d, shape_in=(26, 26, 32), transform=amplitude, filters=32, kernel_size=3, ) x = nengo_dl.tensor_layer(x, nengo_nl, **ensemble_params) x = nengo_dl.tensor_layer( x, tf_compat.layers.average_pooling2d, shape_in=(24, 24, 32), synapse=synapse, transform=amplitude, pool_size=2, strides=2, ) x = nengo_dl.tensor_layer(x, tf_compat.layers.dense, units=128) x = nengo_dl.tensor_layer(x, nengo_nl, **ensemble_params) x = nengo_dl.tensor_layer( x, tf_compat.layers.dropout, rate=0.4, transform=amplitude ) x = nengo_dl.tensor_layer(x, tf_compat.layers.dense, units=10) else: nl = tf.nn.relu # def softlif_layer(x, sigma=1, tau_ref=0.002, tau_rc=0.02, # amplitude=1): # # x -= 1 # z = tf.nn.softplus(x / sigma) * sigma # z += 1e-10 # rates = amplitude / (tau_ref + tau_rc * tf.log1p(1 / z)) # return rates @nengo_dl.reshaped((28, 28, 1)) def mnist_node(_, x): # pragma: no cover x = tf_compat.layers.conv2d(x, filters=32, kernel_size=3, activation=nl) x = tf_compat.layers.conv2d(x, filters=32, kernel_size=3, activation=nl) x = tf_compat.layers.average_pooling2d(x, pool_size=2, strides=2) x = tf_compat.layers.flatten(x) x = tf_compat.layers.dense(x, 128, activation=nl) x = tf_compat.layers.dropout(x, rate=0.4) x = tf_compat.layers.dense(x, 10) return x node = nengo_dl.TensorNode(mnist_node, size_in=28 * 28, size_out=10) x = node nengo.Connection(net.inp, node, synapse=None) net.p = nengo.Probe(x) return net
[docs]def spaun(dimensions): """ Builds the Spaun network from [1]_ Parameters ---------- dimensions : int Number of dimensions for vector values Returns ------- net : `nengo.Network` benchmark network References ---------- .. [1] Chris Eliasmith, Terrence C. Stewart, Xuan Choo, Trevor Bekolay, Travis DeWolf, Yichuan Tang, and Daniel Rasmussen (2012). A large-scale model of the functioning brain. Science, 338:1202-1205. Notes ----- This network needs to be installed via .. code-block:: bash pip install git+https://github.com/drasmuss/spaun2.0.git """ from _spaun.configurator import cfg from _spaun.vocabulator import vocab from _spaun.experimenter import experiment from _spaun.modules.stim import stim_data from _spaun.modules.vision import vis_data from _spaun.modules.motor import mtr_data from _spaun.spaun_main import Spaun vocab.sp_dim = dimensions cfg.mtr_arm_type = None cfg.set_seed(1) experiment.initialize( "A", stim_data.get_image_ind, stim_data.get_image_label, cfg.mtr_est_digit_response_time, "", cfg.rng, ) vocab.initialize(stim_data.stim_SP_labels, experiment.num_learn_actions, cfg.rng) vocab.initialize_mtr_vocab(mtr_data.dimensions, mtr_data.sps) vocab.initialize_vis_vocab(vis_data.dimensions, vis_data.sps) return Spaun()
[docs]def random_network( dimensions, neurons_per_d, neuron_type, n_ensembles, connections_per_ensemble, seed=0, ): """ Basal ganglia network benchmark. Parameters ---------- dimensions : int Number of dimensions for vector values neurons_per_d : int Number of neurons to use per vector dimension neuron_type : `~nengo.neurons.NeuronType` Simulation neuron type n_ensembles : int Number of ensembles in the network connections_per_ensemble : int Outgoing connections from each ensemble Returns ------- net : `nengo.Network` benchmark network """ random.seed(seed) with nengo.Network(label="random", seed=seed) as net: net.inp = nengo.Node([0] * dimensions) net.out = nengo.Node(size_in=dimensions) net.p = nengo.Probe(net.out) ensembles = [ nengo.Ensemble( neurons_per_d * dimensions, dimensions, neuron_type=neuron_type ) for _ in range(n_ensembles) ] dec = np.ones((neurons_per_d * dimensions, dimensions)) for ens in net.ensembles: # add a connection to input and output node, so we never have # any "orphan" ensembles nengo.Connection(net.inp, ens) nengo.Connection(ens, net.out, solver=nengo.solvers.NoSolver(dec)) posts = random.sample(ensembles, connections_per_ensemble) for post in posts: nengo.Connection(ens, post, solver=nengo.solvers.NoSolver(dec)) return net
[docs]def run_profile(net, train=False, n_steps=150, do_profile=True, reps=1, **kwargs): """ Run profiler on a benchmark network. Parameters ---------- net : `~nengo.Network` The nengo Network to be profiled. train : bool If True, profile the ``sim.train`` function. Otherwise, profile the ``sim.run`` function. n_steps : int The number of timesteps to run the simulation. do_profile : bool Whether or not to run profiling reps : int Repeat the run this many times (only profile data from the last run will be kept). Returns ------- exec_time : float Time (in seconds) taken to run the benchmark, taking the minimum over ``reps``. Notes ----- kwargs will be passed on to `.Simulator` """ exec_time = 1e10 with net: nengo_dl.configure_settings(inference_only=not train) with nengo_dl.Simulator(net, **kwargs) as sim: if train: opt = tf_compat.train.GradientDescentOptimizer(0.001) x = np.random.randn(sim.minibatch_size, n_steps, net.inp.size_out) y = np.random.randn(sim.minibatch_size, n_steps, net.p.size_in) # run once to eliminate startup overhead sim.train( {net.inp: x, net.p: y}, optimizer=opt, n_epochs=1, profile=do_profile ) for _ in range(reps): start = time.time() sim.train( {net.inp: x, net.p: y}, optimizer=opt, n_epochs=1, profile=do_profile, ) exec_time = min(time.time() - start, exec_time) print("Execution time:", exec_time) else: # run once to eliminate startup overhead sim.run_steps(n_steps, profile=do_profile) for _ in range(reps): start = time.time() sim.run_steps(n_steps, profile=do_profile) exec_time = min(time.time() - start, exec_time) print("Execution time:", exec_time) return exec_time
@click.group(chain=True) def main(): """Command-line interface for benchmarks.""" @main.command() @click.pass_obj @click.option("--benchmark", default="cconv", help="Name of benchmark network") @click.option("--dimensions", default=128, help="Number of dimensions") @click.option("--neurons_per_d", default=64, help="Neurons per dimension") @click.option("--neuron_type", default="RectifiedLinear", help="Nengo neuron model") @click.option( "--kwarg", type=str, multiple=True, help="Arbitrary kwarg to pass to benchmark network (key=value)", ) def build(obj, benchmark, dimensions, neurons_per_d, neuron_type, kwarg): """Builds one of the benchmark networks""" # get benchmark network by name benchmark = globals()[benchmark] # get the neuron type object from string class name try: neuron_type = getattr(nengo, neuron_type)() except AttributeError: neuron_type = getattr(nengo_dl, neuron_type)() # set up kwargs kwargs = dict((k, int(v)) for k, v in [a.split("=") for a in kwarg]) # add the special cli kwargs if applicable; note we could just do # everything through --kwarg, but it is convenient to have a # direct option for the common arguments params = inspect.signature(benchmark).parameters for kw in ("benchmark", "dimensions", "neurons_per_d", "neuron_type"): if kw in params: kwargs[kw] = locals()[kw] # build benchmark and add to context for chaining print( "Building %s with %s" % (nengo_dl.utils.function_name(benchmark, sanitize=False), kwargs) ) obj["net"] = benchmark(**kwargs) @main.command() @click.pass_obj @click.option( "--train/--no-train", default=False, help="Whether to profile training (as opposed to running) the network", ) @click.option( "--n_steps", default=150, help="Number of steps for which to run the simulation" ) @click.option("--batch_size", default=1, help="Number of inputs to the model") @click.option( "--device", default="/gpu:0", help="TensorFlow device on which to run the simulation", ) @click.option( "--unroll", default=25, help="Number of steps for which to unroll the simulation" ) @click.option( "--time-only", is_flag=True, default=False, help="Only count total time, rather than profiling internals", ) def profile(obj, train, n_steps, batch_size, device, unroll, time_only): """Runs profiling on a network (call after 'build')""" if "net" not in obj: raise ValueError("Must call `build` before `profile`") obj["time"] = run_profile( obj["net"], do_profile=not time_only, train=train, n_steps=n_steps, minibatch_size=batch_size, device=device, unroll_simulation=unroll, ) @main.command() def matmul_vs_reduce(): # pragma: no cover """ Compares two different approaches to batched matrix multiplication (tf.matmul vs tf.multiply+tf.reduce_sum). This is relevant for figuring out which approach is more efficient on a given system for different matrix shapes (determining which method we use in DotIncBuilder). """ # a_shape = (n_ops, s0, s1, 1) # x_shape = (n_ops, 1, s1, mini) # for matmul we omit the 1 dimensions a_c = tf_compat.placeholder(tf.float64, shape=(None, None, None), name="a_c") x_c = tf_compat.placeholder(tf.float64, shape=(None, None, None), name="b_c") a_d = tf_compat.placeholder(tf.float64, shape=(None, None, None, 1), name="a_d") x_d = tf_compat.placeholder(tf.float64, shape=(None, 1, None, None), name="b_d") c = tf.matmul(a_c, x_c) d = tf.reduce_sum(input_tensor=tf.multiply(a_d, x_d), axis=-2) reps = 100 n_ops_range = [1, 4, 8, 16, 32, 64] mini_range = [1, 16, 32, 64, 128] s0_range = [1, 64, 128, 192, 256] s1_range = [1, 64, 128, 192, 256] matmul_times = np.zeros( (len(n_ops_range), len(mini_range), len(s0_range), len(s1_range)) ) reduce_times = np.zeros_like(matmul_times) params = itertools.product( enumerate(n_ops_range), enumerate(mini_range), enumerate(s0_range), enumerate(s1_range), ) with tf_compat.Session() as sess: for (i, n_ops), (j, mini), (k, s0), (l, s1) in params: print(n_ops, mini, s0, s1) a_val = np.random.randn(n_ops, s0, s1, 1) x_val = np.random.randn(n_ops, 1, s1, mini) for r in range(reps + 3): if r == 3: start = time.time() c_val = sess.run(c, feed_dict={a_c: a_val[..., 0], x_c: x_val[:, 0]}) matmul_times[i, j, k, l] = (time.time() - start) / reps for r in range(reps + 3): if r == 3: start = time.time() d_val = sess.run(d, feed_dict={a_d: a_val, x_d: x_val}) reduce_times[i, j, k, l] = (time.time() - start) / reps assert np.allclose(c_val, d_val) fig, ax = plt.subplots(len(n_ops_range), len(mini_range), sharex=True, sharey=True) X, Y = np.meshgrid(s0_range, s1_range) Z = matmul_times - reduce_times v = np.sort(np.concatenate((np.linspace(np.min(Z), np.max(Z), 10), [0]))) for i, n_ops in enumerate(n_ops_range): for j, mini in enumerate(mini_range): cs = ax[i][j].contourf(X, Y, Z[i, j], v) if i == 0: ax[i][j].set_title("mini %d" % mini) if j == 0: ax[i][j].set_ylabel("ops %d" % n_ops) DATA_DIR = os.path.join(os.path.dirname(nengo_dl.__file__), "..", "data") np.savez( os.path.join(DATA_DIR, "matmul_benchmarks"), n_ops_range, mini_range, s0_range, s1_range, Z, ) fig.subplots_adjust(right=0.8) cbar_ax = fig.add_axes([0.85, 0.15, 0.05, 0.7]) fig.colorbar(cs, cax=cbar_ax) plt.show() if __name__ == "__main__": main(obj={}) # pragma: no cover