def test_identity_array(n_ensembles, ens_dimensions):
with nengo.Network() as model:
a = nengo.networks.EnsembleArray(10, n_ensembles, ens_dimensions)
b = nengo.networks.EnsembleArray(10, n_ensembles, ens_dimensions)
nengo.Connection(a.output, b.input)
nengo_loihi.add_params(model)
networks = splitter.split(model,
precompute=False,
remove_passthrough=True,
max_rate=1000,
inter_tau=0.005)
# ignore the a.input -> a.ensemble connections
connections = [
c for c in networks.chip.connections
if not (isinstance(c.pre_obj, splitter.ChipReceiveNode)
and c.post_obj in a.ensembles)
]
assert len(connections) == n_ensembles
pre = set()
post = set()
for c in connections:
assert c.pre in a.all_ensembles or c.pre_obj is a.input
assert c.post in b.all_ensembles
assert np.allclose(c.transform, np.eye(ens_dimensions))
pre.add(c.pre)
post.add(c.post)
assert len(pre) == n_ensembles
assert len(post) == n_ensembles
def test_probedict_fallbacks(precompute, Simulator):
with nengo.Network() as net:
nengo_loihi.add_params(net)
node_a = nengo.Node(0)
with nengo.Network():
ens_b = nengo.Ensemble(10, 1)
conn_ab = nengo.Connection(node_a, ens_b)
ens_c = nengo.Ensemble(5, 1)
net.config[ens_c].on_chip = False
conn_bc = nengo.Connection(ens_b, ens_c)
probe_a = nengo.Probe(node_a)
probe_c = nengo.Probe(ens_c)
with Simulator(net, precompute=precompute) as sim:
sim.run(0.002)
assert node_a in sim.data
assert ens_b in sim.data
assert ens_c in sim.data
assert probe_a in sim.data
assert probe_c in sim.data
# TODO: connections are currently not probeable as they are
# replaced in the splitting process
assert conn_ab # in sim.data
assert conn_bc # in sim.data
def test_full_array(n_ensembles, ens_dimensions):
with nengo.Network() as model:
a = nengo.networks.EnsembleArray(10, n_ensembles, ens_dimensions)
b = nengo.networks.EnsembleArray(10, n_ensembles, ens_dimensions)
D = n_ensembles * ens_dimensions
nengo.Connection(a.output, b.input, transform=np.ones((D, D)))
nengo_loihi.add_params(model)
networks = splitter.split(model,
precompute=False,
remove_passthrough=True,
max_rate=1000,
inter_tau=0.005)
# ignore the a.input -> a.ensemble connections
connections = [
c for c in networks.chip.connections
if not (isinstance(c.pre_obj, splitter.ChipReceiveNode)
and c.post_obj in a.ensembles)
]
assert len(connections) == n_ensembles**2
pairs = set()
for c in connections:
assert c.pre in a.all_ensembles
assert c.post in b.all_ensembles
assert np.allclose(c.transform,
np.ones((ens_dimensions, ens_dimensions)))
pairs.add((c.pre, c.post))
assert len(pairs) == n_ensembles**2
def test_sliced_probe(allclose, probe_slice, Simulator):
n_neurons = 16
# The bias should be the same for each block, but different within the block
# to see some variety. This gives bias 0, 1, 2, 3 to the four neurons in the
# 2 by 2 block.
bias = np.tile(np.arange(4).reshape((2, 2)), (2, 2)).flatten()
with nengo.Network() as net:
e = nengo.Ensemble(n_neurons, 1, gain=np.zeros(n_neurons), bias=bias)
p = nengo.Probe(e.neurons[probe_slice], "voltage", synapse=0.002)
with Simulator(net) as ref_sim:
ref_sim.run(0.01)
ref_voltages = ref_sim.data[p]
with net:
nengo_loihi.add_params(net)
net.config[e].block_shape = nengo_loihi.BlockShape((2, 2), (n_neurons // 4, 4))
with Simulator(net) as split_sim:
split_sim.run(ref_sim.time)
split_voltages = split_sim.data[p]
assert np.all(ref_sim.data[e].gain == split_sim.data[e].gain)
assert np.all(ref_sim.data[e].bias == split_sim.data[e].bias)
assert allclose(split_voltages, ref_voltages)
def test_ens_decoded_on_host(precompute, allclose, Simulator, seed, plt):
out_synapse = nengo.synapses.Alpha(0.03)
simtime = 0.6
with nengo.Network(seed=seed) as model:
nengo_loihi.add_params(model)
stim = nengo.Node(lambda t: [np.sin(t * 2 * np.pi / simtime)])
a = nengo.Ensemble(100, 1)
model.config[a].on_chip = False
b = nengo.Ensemble(100, 1)
nengo.Connection(stim, a)
nengo.Connection(a, b, function=lambda x: -x)
p_stim = nengo.Probe(stim, synapse=out_synapse)
p_a = nengo.Probe(a, synapse=out_synapse)
p_b = nengo.Probe(b, synapse=out_synapse)
with Simulator(model, precompute=precompute) as sim:
sim.run(simtime)
plt.plot(sim.trange(), sim.data[p_stim])
plt.plot(sim.trange(), sim.data[p_a])
plt.plot(sim.trange(), sim.data[p_b])
assert allclose(sim.data[p_a], sim.data[p_stim], atol=0.05, rtol=0.01)
assert allclose(sim.data[p_b], -sim.data[p_a], atol=0.15, rtol=0.1)
def test_probe_split_blocks(Simulator, seed, plt):
n_neurons = 80
gain = np.ones(n_neurons)
bias = np.linspace(0, 20, n_neurons)
simtime = 0.2
with nengo.Network(seed=seed) as net:
ens = nengo.Ensemble(n_neurons, 1, gain=gain, bias=bias)
probe = nengo.Probe(ens.neurons)
probe1_slice = slice(3, 33)
probe1 = nengo.Probe(ens.neurons[probe1_slice])
probe2_slice = slice(7, 52, 3)
probe2 = nengo.Probe(ens.neurons[probe2_slice])
probe3_slice = [2, 5, 17, 21, 36, 49, 52, 69,
73] # randomly chosen inds
probe3 = nengo.Probe(ens.neurons[probe3_slice])
# run without splitting ensemble
with Simulator(net) as sim1:
assert len(sim1.model.blocks) == 1
sim1.run(simtime)
# run with splitting ensemble
with net:
add_params(net)
net.config[ens].block_shape = BlockShape((5, 4), (10, 8))
with Simulator(net) as sim2:
assert len(sim2.model.blocks) == 4
sim2.run(simtime)
for k, sim in enumerate((sim1, sim2)):
plt.subplot(2, 1, k + 1)
plt.plot(bias, sim.data[probe].mean(axis=0))
plt.plot(bias[probe1_slice], sim.data[probe1].mean(axis=0))
plt.plot(bias[probe2_slice], sim.data[probe2].mean(axis=0), ".")
plt.plot(bias[probe3_slice], sim.data[probe3].mean(axis=0), "x")
# ensure rates increase and not everything is zero
for sim in (sim1, sim2):
diffs = np.diff(sim.data[probe].mean(axis=0))
assert (diffs >= 0).all() and (diffs > 1).sum() > 10
# ensure slices match unsliced probe
for sim in (sim1, sim2):
assert np.array_equal(sim.data[probe1], sim.data[probe][:,
probe1_slice])
assert np.array_equal(sim.data[probe2], sim.data[probe][:,
probe2_slice])
assert np.array_equal(sim.data[probe3], sim.data[probe][:,
probe3_slice])
# ensure split and unsplit simulators match
for p in (probe, probe1, probe2, probe3):
assert np.array_equal(sim1.data[p], sim2.data[p])
def test_model_validate_notempty(Simulator):
with nengo.Network() as model:
nengo_loihi.add_params(model)
a = nengo.Ensemble(10, 1)
model.config[a].on_chip = False
with pytest.raises(nengo.exceptions.BuildError):
with Simulator(model):
pass
def test_dt(dt, pre_on_chip, Simulator, seed, plt, allclose):
function = lambda x: -x
simtime = 0.2
probe_synapse = nengo.Alpha(0.01)
conn_synapse = nengo.Lowpass(0.005)
stim_synapse = probe_synapse
# stim synapse accounts for delays in connections/probes, so we can compare
if pre_on_chip:
stim_synapse = stim_synapse.combine(conn_synapse)
ens_params = dict(
intercepts=nengo.dists.Uniform(-0.9, 0.9),
max_rates=nengo.dists.Uniform(100, 120),
)
with nengo.Network(seed=seed) as model:
nengo_loihi.add_params(model)
stim = nengo.Node(lambda t: -(np.sin(2 * np.pi * t / simtime)))
stim_p = nengo.Probe(stim, synapse=stim_synapse)
pre = nengo.Ensemble(100, 1, **ens_params)
model.config[pre].on_chip = pre_on_chip
pre_p = nengo.Probe(pre, synapse=probe_synapse)
post = nengo.Ensemble(101, 1, **ens_params)
post_p = nengo.Probe(post, synapse=probe_synapse)
nengo.Connection(stim, pre, synapse=None)
nengo.Connection(
pre,
post,
function=function,
synapse=conn_synapse,
solver=nengo.solvers.LstsqL2(weights=True),
)
with Simulator(model, dt=dt) as sim:
assert sim.model.decode_tau == conn_synapse.tau
sim.run(simtime)
x = sim.data[stim_p]
y = function(x)
if pre_on_chip:
y = conn_synapse.filt(y, dt=dt)
else:
y = conn_synapse.combine(conn_synapse).filt(y, dt=dt)
plt.plot(sim.trange(), x, "k--")
plt.plot(sim.trange(), y, "k--")
plt.plot(sim.trange(), sim.data[pre_p])
plt.plot(sim.trange(), sim.data[post_p])
assert allclose(sim.data[pre_p], x, rtol=0.1, atol=0.1)
assert allclose(sim.data[post_p], y, rtol=0.1, atol=0.1)
def test_n2n_on_host(precompute, allclose, Simulator, seed_ens, seed, plt):
"""Ensure that neuron to neuron connections work on and off chip."""
n_neurons = 50
# When the ensemble is seeded, the output plots will make more sense,
# but the test should work whether they're seeded or not.
ens_seed = (seed + 1) if seed_ens else None
if not seed_ens:
pytest.xfail("Seeds change when moving ensembles off/on chip")
simtime = 1.0
with nengo.Network(seed=seed) as model:
nengo_loihi.add_params(model)
stim = nengo.Node(lambda t: [np.sin(t * 2 * np.pi / simtime)])
# pre receives stimulation and represents the sine wave
pre = nengo.Ensemble(n_neurons, dimensions=1, seed=ens_seed)
model.config[pre].on_chip = False
nengo.Connection(stim, pre)
# post has pre's neural activity forwarded to it.
# Since the neuron parameters are the same, it should also represent
# the same sine wave.
# The 0.015 scaling is chosen so the values match visually,
# though a more principled reason would be better.
post = nengo.Ensemble(n_neurons, dimensions=1, seed=ens_seed)
nengo.Connection(pre.neurons,
post.neurons,
transform=np.eye(n_neurons) * 0.015)
p_synapse = nengo.synapses.Alpha(0.03)
p_stim = nengo.Probe(stim, synapse=p_synapse)
p_pre = nengo.Probe(pre, synapse=p_synapse)
p_post = nengo.Probe(post, synapse=p_synapse)
with Simulator(model, precompute=precompute) as sim:
sim.run(simtime)
t = sim.trange()
model.config[pre].on_chip = True
with Simulator(model, precompute=precompute) as sim2:
sim2.run(simtime)
t2 = sim2.trange()
plt.plot(t, sim.data[p_stim], c="k", label="input")
plt.plot(t, sim.data[p_pre], label="pre off-chip")
plt.plot(t, sim.data[p_post], label="post (pre off-chip)")
plt.plot(t2, sim2.data[p_pre], label="pre on-chip")
plt.plot(t2, sim2.data[p_post], label="post (pre on-chip)")
plt.legend()
assert allclose(sim.data[p_pre], sim2.data[p_pre], atol=0.1)
assert allclose(sim.data[p_post], sim2.data[p_post], atol=0.1)
def test_conv_input(channels_last, Simulator, plt, allclose):
input_shape = ImageShape(4, 4, 1, channels_last=channels_last)
seed = 3 # fix seed to do the same computation for both channel positions
rng = np.random.RandomState(seed + 1)
with nengo.Network(seed=seed) as net:
nengo_loihi.add_params(net)
a = nengo.Node(rng.uniform(0, 1, size=input_shape.size))
nc = 2
kernel = np.array([1., -1.]).reshape((1, 1, 1, nc))
transform = nengo_loihi.Conv2D.from_kernel(kernel, input_shape)
b = nengo.Ensemble(transform.output_shape.size,
1,
neuron_type=nengo.SpikingRectifiedLinear(),
max_rates=nengo.dists.Choice([50]),
intercepts=nengo.dists.Choice([0]))
net.config[b].on_chip = False
nengo.Connection(a, b.neurons, transform=transform)
output_shape = transform.output_shape
nf = 4
kernel = rng.uniform(-0.005, 0.005, size=(nc, 3, 3, nf))
transform = nengo_loihi.Conv2D.from_kernel(kernel, output_shape)
c = nengo.Ensemble(transform.output_shape.size,
1,
neuron_type=nengo.LIF(),
max_rates=nengo.dists.Choice([100]),
intercepts=nengo.dists.Choice([0]))
nengo.Connection(b.neurons, c.neurons, transform=transform)
output_shape = transform.output_shape
p = nengo.Probe(c.neurons)
with nengo.Simulator(net, optimize=False) as sim:
sim.run(1.0)
with Simulator(net, seed=seed) as sim_loihi:
sim_loihi.run(1.0)
p0 = np.sum(sim.data[p] > 0, axis=0).reshape(output_shape.shape())
p1 = np.sum(sim_loihi.data[p] > 0, axis=0).reshape(output_shape.shape())
if not output_shape.channels_last:
p0 = np.transpose(p0, (1, 2, 0))
p1 = np.transpose(p1, (1, 2, 0))
plt.plot(p0.ravel(), 'k')
plt.plot(p1.ravel(), 'b--')
# loihi spikes are not exactly the same, but should be close-ish
assert allclose(p0, p1, rtol=0.15, atol=1)
def test_model_validate_notempty(Simulator):
with nengo.Network() as model:
nengo_loihi.add_params(model)
a = nengo.Ensemble(10, 1)
model.config[a].on_chip = False
assert nengo.rc.get("decoder_cache", "enabled")
with pytest.raises(BuildError, match="No neurons marked"):
with Simulator(model):
pass
# Ensure cache config not changed
assert nengo.rc.get("decoder_cache", "enabled")
def piecewise_net(n_pres, pres_time, seed):
values = np.linspace(-1, 1, n_pres)
with nengo.Network(seed=seed) as net:
add_params(net)
inp = nengo.Node(nengo.processes.PresentInput(values, pres_time),
size_out=1)
ens = nengo.Ensemble(100, 1)
nengo.Connection(inp, ens)
net.probe = nengo.Probe(ens, synapse=nengo.Alpha(0.01))
node = nengo.Node(size_in=1)
nengo.Connection(ens, node, synapse=nengo.Alpha(0.01))
net.node_probe = nengo.Probe(node)
return net, values
def test_interchip_helpers(Simulator, rng):
"""Test other cases for helper functions used by Interchip allocators."""
with nengo.Network() as net:
nengo_loihi.add_params(net)
a = nengo.Ensemble(2000, 1) # this ensemble will be split (2 blocks)
b = nengo.Ensemble(1000, 1) # this ensemble will be one block
c = nengo.Ensemble(1000, 1) # this ensemble will be off-chip
net.config[c].on_chip = False
nengo.Connection(a, b)
with nengo.Network():
d = nengo.Ensemble(10, 1) # this ensemble is in a different network (errors)
ens_rates = {
a: rng.uniform(size=a.n_neurons),
b: rng.uniform(size=b.n_neurons),
c: rng.uniform(size=c.n_neurons),
}
with Simulator(net) as sim:
# --- test ensemble_to_block_rates
block_rates = ensemble_to_block_rates(sim.model, ens_rates)
a_blocks = sim.model.objs[a]["out"]
b_blocks = sim.model.objs[b]["out"]
assert set(block_rates) == (set(a_blocks) | set(b_blocks))
i = 0
for block in a_blocks:
assert np.array_equal(
ens_rates[a][i : i + block.n_neurons], block_rates[block]
)
i += block.n_neurons
assert len(b_blocks) == 1 and np.array_equal(
ens_rates[b], block_rates[b_blocks[0]]
)
# test ValueError if ensemble not in model
ens_rates[d] = rng.uniform(size=d.n_neurons)
with pytest.raises(ValueError, match="Ensemble.*does not appear in the model"):
ensemble_to_block_rates(sim.model, ens_rates)
# --- test estimate_interblock_activity error
partial_ens_rates = {a: ens_rates[a]}
with pytest.raises(KeyError, match="block.*not in block_rates"):
GreedyInterchip(ensemble_rates=partial_ens_rates)(sim.model, n_chips=2)
def test_synapse_merging(Simulator, seed):
with nengo.Network(seed=seed) as model:
a = nengo.networks.EnsembleArray(10, n_ensembles=2)
b = nengo.Node(None, size_in=2)
c = nengo.networks.EnsembleArray(10, n_ensembles=2)
nengo.Connection(a.output[0], b[0], synapse=None)
nengo.Connection(a.output[1], b[1], synapse=0.1)
nengo.Connection(b[0], c.input[0], synapse=None)
nengo.Connection(b[0], c.input[1], synapse=0.2)
nengo.Connection(b[1], c.input[0], synapse=None)
nengo.Connection(b[1], c.input[1], synapse=0.2)
nengo_loihi.add_params(model)
networks = splitter.split(model,
precompute=False,
remove_passthrough=True,
max_rate=1000,
inter_tau=0.005)
# ignore the a.input -> a.ensemble connections
connections = [
c for c in networks.chip.connections
if not (isinstance(c.pre_obj, splitter.ChipReceiveNode)
and c.post_obj in a.ensembles)
]
assert len(connections) == 4
desired_filters = {
('0', '0'): None,
('0', '1'): 0.2,
('1', '0'): 0.1,
('1', '1'): 0.3,
}
for c in connections:
if desired_filters[(c.pre.label, c.post.label)] is None:
assert c.synapse is None
else:
assert isinstance(c.synapse, nengo.Lowpass)
assert np.allclose(c.synapse.tau,
desired_filters[(c.pre.label, c.post.label)])
# check that model builds/runs correctly
with Simulator(model, remove_passthrough=True) as sim:
sim.step()
def test_seeds(precompute, Simulator, seed):
with nengo.Network(seed=seed) as net:
nengo_loihi.add_params(net)
e0 = nengo.Ensemble(1, 1, label="e0")
e1 = nengo.Ensemble(1, 1, seed=2, label="e1")
e2 = nengo.Ensemble(1, 1, label="e2")
net.config[e2].on_chip = False
nengo.Connection(e0, e1)
nengo.Connection(e0, e2)
with nengo.Network():
n = nengo.Node(0)
e = nengo.Ensemble(1, 1, label="e")
nengo.Node(1)
nengo.Connection(n, e)
nengo.Probe(e)
with nengo.Network(seed=8):
nengo.Ensemble(8, 1, seed=3, label="unnamed")
nengo.Node(1)
def get_seed(sim, obj):
return sim.model.seeds.get(
obj,
sim.model.host.seeds.get(obj,
sim.model.host_pre.seeds.get(obj, None)))
# --- test that seeds are the same as nengo ref simulator
ref = nengo.Simulator(net)
with Simulator(net, precompute=precompute) as sim:
for obj in net.all_objects:
assert get_seed(sim, obj) == ref.model.seeds.get(obj, None)
# --- test that seeds that we set are preserved after splitting
model = nengo_loihi.builder.Model()
for i, obj in enumerate(net.all_objects):
model.seeds[obj] = i
with Simulator(net, model=model, precompute=precompute) as sim:
for i, obj in enumerate(net.all_objects):
assert get_seed(sim, obj) == i
def test_dt(dt, pre_on_chip, Simulator, seed, plt, allclose):
function = lambda x: x**2
probe_synapse = nengo.Alpha(0.01)
simtime = 0.2
ens_params = dict(
intercepts=nengo.dists.Uniform(-0.9, 0.9),
max_rates=nengo.dists.Uniform(100, 120),
)
with nengo.Network(seed=seed) as model:
nengo_loihi.add_params(model)
stim = nengo.Node(lambda t: -(np.sin(2 * np.pi * t / simtime)))
stim_p = nengo.Probe(stim, synapse=probe_synapse)
pre = nengo.Ensemble(100, 1, **ens_params)
model.config[pre].on_chip = pre_on_chip
pre_p = nengo.Probe(pre, synapse=probe_synapse)
post = nengo.Ensemble(101, 1, **ens_params)
post_p = nengo.Probe(post, synapse=probe_synapse)
nengo.Connection(stim, pre)
nengo.Connection(pre,
post,
function=function,
solver=nengo.solvers.LstsqL2(weights=True))
with Simulator(model, dt=dt) as sim:
sim.run(simtime)
x = sim.data[stim_p]
y = function(x)
plt.plot(sim.trange(), x, "k--")
plt.plot(sim.trange(), y, "k--")
plt.plot(sim.trange(), sim.data[pre_p])
plt.plot(sim.trange(), sim.data[post_p])
assert allclose(sim.data[pre_p], x, rtol=0.1, atol=0.1)
assert allclose(sim.data[post_p], y, rtol=0.1, atol=0.1)
def __init__(self, input_node, n_neuron_out, lr, onchip=True):
'''
initialize critic net as a nengo network object
PARAMS:
n_pc - number of place cells
n_neuron_in - number of neurons in Ensemble encoding input
n_neuron_out - number of neurons in Ensemble encoding output
'''
with nengo.Network() as net:
nengo_loihi.add_params(net)
net.output = nengo.Ensemble(n_neurons=n_neuron_out, dimensions=1)
net.conn = nengo.Connection(input_node,
net.output,
function=lambda x: [0],
synapse=0.01)
# TODO: PES changes Decoders of incoming node
# TODO: Does this interfer with other learning due to shared input?
net.conn.learning_rule_type = nengo.PES(learning_rate=lr)
net.config[net.output].on_chip = onchip
self.net = net
def test_slicing_bugs(Simulator, seed):
n = 50
with nengo.Network() as model:
a = nengo.Ensemble(n, 1, label="a")
p0 = nengo.Probe(a[0])
p = nengo.Probe(a)
with Simulator(model) as sim:
sim.run(0.1)
assert np.allclose(sim.data[p0], sim.data[p])
assert a in sim.model.params
assert a not in sim.model.host.params
with nengo.Network() as model:
nengo_loihi.add_params(model)
a = nengo.Ensemble(n, 1, label="a")
b0 = nengo.Ensemble(n, 1, label="b0", seed=seed)
model.config[b0].on_chip = False
nengo.Connection(a[0], b0)
b = nengo.Ensemble(n, 1, label="b", seed=seed)
model.config[b].on_chip = False
nengo.Connection(a, b)
p0 = nengo.Probe(b0)
p = nengo.Probe(b)
with Simulator(model) as sim:
sim.run(0.1)
assert np.allclose(sim.data[p0], sim.data[p])
assert a in sim.model.params
assert a not in sim.model.host.params
assert b not in sim.model.params
assert b in sim.model.host.params
def test_passthrough_placement():
with nengo.Network() as model:
stim = nengo.Node(0)
a = nengo.Node(None, size_in=1) # should be off-chip
b = nengo.Ensemble(10, 1)
c = nengo.Node(None, size_in=1) # should be removed
d = nengo.Node(None, size_in=1) # should be removed
e = nengo.Node(None, size_in=1) # should be removed
f = nengo.Ensemble(10, 1)
g = nengo.Node(None, size_in=1) # should be off-chip
nengo.Connection(stim, a)
nengo.Connection(a, b)
nengo.Connection(b, c)
nengo.Connection(c, d)
nengo.Connection(d, e)
nengo.Connection(e, f)
nengo.Connection(f, g)
nengo.Probe(g)
nengo_loihi.add_params(model)
networks = splitter.split(model,
precompute=False,
remove_passthrough=True,
max_rate=1000,
inter_tau=0.005)
chip = networks.chip
host = networks.host
assert a in host.nodes
assert a not in chip.nodes
assert c not in host.nodes
assert c not in chip.nodes
assert d not in host.nodes
assert d not in chip.nodes
assert e not in host.nodes
assert e not in chip.nodes
assert g in host.nodes
assert g not in chip.nodes
def test_transform_merging(d1, d2, d3):
with nengo.Network() as model:
a = nengo.Ensemble(10, d1)
b = nengo.Node(None, size_in=d2)
c = nengo.Ensemble(10, d3)
t1 = np.random.uniform(-1, 1, (d2, d1))
t2 = np.random.uniform(-1, 1, (d3, d2))
nengo.Connection(a, b, transform=t1)
nengo.Connection(b, c, transform=t2)
nengo_loihi.add_params(model)
networks = splitter.split(model,
precompute=False,
remove_passthrough=True,
max_rate=1000,
inter_tau=0.005)
chip = networks.chip
assert len(chip.connections) == 1
conn = chip.connections[0]
assert np.allclose(conn.transform, np.dot(t2, t1))
def test_conv_deepnet(
channels_last,
pop_type,
precompute,
Simulator,
request,
rng,
seed,
plt,
allclose,
):
"""Run a convolutional network with two layers on the chip.
Checks that network with block splitting on the target matches one without
on the emulator.
"""
# TODO: This case fails in NxSDK 0.9.0 but will be fixed in the next version.
# Remove this check once the next version is released.
if pop_type == 32:
pytest.skip("Pop32 multichip test requires latest NxSDK")
def set_partition(partition):
os.environ["PARTITION"] = partition
request.addfinalizer(lambda: set_partition(""))
# multichip pop_type = 16 works only on nahuku32 board currently
if pop_type == 16:
set_partition("nahuku32")
def conv_layer(x,
input_shape,
array_init=None,
label=None,
conn_args=None,
**conv_args):
conn_args = {} if conn_args is None else conn_args
if array_init is not None:
assert all(a not in conv_args
for a in ("init", "kernel_size", "n_filters"))
assert array_init.ndim == 4
conv_args["init"] = array_init
conv_args["kernel_size"] = array_init.shape[:2]
assert array_init.shape[2] == input_shape.n_channels
conv_args["n_filters"] = array_init.shape[3]
conv = nengo.Convolution(input_shape=input_shape, **conv_args)
# add an ensemble to implement the activation function
layer = nengo.Ensemble(conv.output_shape.size, 1, label=label)
# connect up the input object to the new layer
conn = nengo.Connection(x, layer.neurons, transform=conv)
return layer, conv, conn
channels = 1
n_filters0 = 1
n_filters1 = 4
n_filters2 = 4
# load data
with open(os.path.join(test_dir, "mnist10.pkl"), "rb") as f:
test10 = pickle.load(f)
test_x = test10[0][0].reshape(28, 28) # range (0, 1)
input_shape = make_channel_shape(test_x.shape, channels, channels_last)
filters0 = np.ones((1, 1, channels, n_filters0))
# use Gabor filters for first layer
filters1 = Gabor(freq=Uniform(0.5, 1),
sigma_x=Choice([0.9]),
sigma_y=Choice([0.9])).generate(n_filters1, (7, 7),
rng=rng)
assert n_filters0 == 1
filters1 = filters1[None, :, :, :] # single channel
filters1 = np.transpose(filters1,
(2, 3, 0, 1)) # rows, cols, in_chan, out_chan
# use random combinations of first-layer channels in 1x1 convolution
filters2 = rng.uniform(-0.2, 1,
size=(n_filters1, n_filters2)).clip(0, None)
filters2 *= 2 / filters2.sum(axis=0,
keepdims=True) # each filter sums to 2
filters2 = filters2[None, None, :, :] # rows, cols, in_chan, out_chan
tau_s = 0.001
max_rate = 100
amp = 1 / max_rate
f_split = 2 if pop_type == 32 else 4
# use Loihi neuron type so Nengo sim mimics Loihi neuron effects
neuron_type = LoihiSpikingRectifiedLinear(amplitude=amp)
pres_time = 0.2
with nengo.Network(seed=seed) as net:
nengo_loihi.add_params(net)
net.config[nengo.Ensemble].neuron_type = neuron_type
net.config[nengo.Ensemble].max_rates = Choice([max_rate])
net.config[nengo.Ensemble].intercepts = Choice([0])
net.config[nengo.Connection].synapse = tau_s
u = nengo.Node(test_x.ravel(), label="u")
layer0, conv0, conn0 = conv_layer(
u,
input_shape=input_shape,
array_init=filters0,
strides=(1, 1),
channels_last=channels_last,
label="layer0",
conn_args=dict(synapse=None),
)
net.config[layer0].on_chip = False
layer1, conv1, conn1 = conv_layer(
layer0.neurons,
input_shape=conv0.output_shape,
array_init=filters1,
strides=(2, 2),
channels_last=channels_last,
label="layer1",
)
net.config[layer1].block_shape = nengo_loihi.BlockShape(
make_shape((4, 4), f_split, channels_last), conv1)
net.config[conn1].pop_type = pop_type
layer2, conv2, conn2 = conv_layer(
layer1.neurons,
input_shape=conv1.output_shape,
array_init=filters2,
strides=(1, 1),
channels_last=channels_last,
label="layer2",
)
net.config[layer2].block_shape = nengo_loihi.BlockShape(
make_shape((4, 4), f_split, channels_last), conv2)
net.config[conn2].pop_type = pop_type
output_p = nengo.Probe(layer2.neurons)
output_shape = conv2.output_shape
with nengo.Simulator(net, optimize=False) as sim_nengo:
sim_nengo.run(pres_time)
ref_out = (sim_nengo.data[output_p] > 0).sum(axis=0).reshape(
output_shape.shape)
with Simulator(net, target="sim") as sim_emu:
sim_emu.run(pres_time)
emu_out = (sim_emu.data[output_p] > 0).sum(axis=0).reshape(
output_shape.shape)
# TODO: Remove the if condition when configurable timeout parameter
# is available in nxsdk
if (pop_type == 32 or
os.popen("sinfo -h --partition=nahuku32").read().find("idle") > 0):
with Simulator(
net,
precompute=precompute,
hardware_options={
"allocator": RoundRobin(),
"snip_max_spikes_per_step": 800,
},
) as sim_loihi:
sim_loihi.run(pres_time)
sim_out = ((sim_loihi.data[output_p] > 0).sum(axis=0).reshape(
output_shape.shape))
elif nengo_loihi.version.dev is None:
pytest.fail(
"Pop16 multichip test failed since Nahuku32 is unavailable")
else:
pytest.skip(
"Pop16 multichip test skipped since Nahuku32 is unavailable")
out_max = ref_out.max()
ref_out = ref_out / out_max
emu_out = emu_out / out_max
sim_out = sim_out / out_max
if channels_last:
# channels first, to display channels in separate plots
ref_out = np.transpose(ref_out, (2, 0, 1))
emu_out = np.transpose(emu_out, (2, 0, 1))
sim_out = np.transpose(sim_out, (2, 0, 1))
# --- plot results
rows = 2
cols = 3
ax = plt.subplot(rows, cols, 1)
imshow(test_x, vmin=0, vmax=1, ax=ax)
ax = plt.subplot(rows, cols, 2)
tile(np.transpose(filters1, (2, 3, 0, 1))[0],
rows=2,
cols=2,
grid=True,
ax=ax)
ax = plt.subplot(rows, cols, 3)
plt.hist((ref_out.ravel(), emu_out.ravel(), sim_out.ravel()), bins=21)
ax = plt.subplot(rows, cols, 4)
tile(ref_out, rows=2, cols=2, grid=True, ax=ax)
ax = plt.subplot(rows, cols, 5)
tile(emu_out, rows=2, cols=2, grid=True, ax=ax)
ax = plt.subplot(rows, cols, 6)
tile(sim_out, rows=2, cols=2, grid=True, ax=ax)
assert allclose(sim_out, ref_out, atol=0.15, rtol=1e-3)
assert allclose(sim_out, emu_out, atol=1e-3, rtol=1e-3)
def test_conv_connection(channels, Simulator, seed, rng, plt, allclose):
# channels_last = True
channels_last = False
if channels > 1:
pytest.xfail("Cannot send population spikes to chip")
# load data
with open(os.path.join(test_dir, 'mnist10.pkl'), 'rb') as f:
test10 = pickle.load(f)
test_x = test10[0][0].reshape(28, 28)
test_x = 1.999 * test_x - 0.999 # range (-1, 1)
test_x = test_x[:, :, None] # single channel
input_shape = ImageShape(test_x.shape[0],
test_x.shape[1],
channels,
channels_last=channels_last)
filters = Gabor(freq=Uniform(0.5, 1)).generate(8, (7, 7), rng=rng)
filters = filters[None, :, :, :] # single channel
filters = np.transpose(filters, (0, 2, 3, 1)) # filters last
strides = (2, 2)
tau_rc = 0.02
tau_ref = 0.002
tau_s = 0.005
dt = 0.001
neuron_type = LoihiLIF(tau_rc=tau_rc, tau_ref=tau_ref)
pres_time = 0.1
with nengo.Network(seed=seed) as model:
nengo_loihi.add_params(model)
u = nengo.Node(nengo.processes.PresentInput([test_x.ravel()],
pres_time),
label='u')
a = nengo.Ensemble(input_shape.size,
1,
neuron_type=LoihiSpikingRectifiedLinear(),
max_rates=nengo.dists.Choice([40 / channels]),
intercepts=nengo.dists.Choice([0]),
label='a')
model.config[a].on_chip = False
if channels == 1:
nengo.Connection(u, a.neurons, transform=1, synapse=None)
elif channels == 2:
# encode image into spikes using two channels (on/off)
if input_shape.channels_last:
nengo.Connection(u, a.neurons[0::2], transform=1, synapse=None)
nengo.Connection(u,
a.neurons[1::2],
transform=-1,
synapse=None)
else:
k = input_shape.rows * input_shape.cols
nengo.Connection(u, a.neurons[:k], transform=1, synapse=None)
nengo.Connection(u, a.neurons[k:], transform=-1, synapse=None)
filters = np.vstack([filters, -filters])
else:
raise ValueError("Test not configured for more than two channels")
conv2d_transform = Conv2D.from_kernel(filters,
input_shape,
strides=strides)
output_shape = conv2d_transform.output_shape
gain, bias = neuron_type.gain_bias(max_rates=100, intercepts=0)
gain = gain * 0.01 # account for `a` max_rates
b = nengo.Ensemble(output_shape.size,
1,
neuron_type=neuron_type,
gain=nengo.dists.Choice([gain[0]]),
bias=nengo.dists.Choice([bias[0]]),
label='b')
nengo.Connection(a.neurons,
b.neurons,
synapse=tau_s,
transform=conv2d_transform)
bp = nengo.Probe(b.neurons)
with nengo.Simulator(model, dt=dt, optimize=False) as sim:
sim.run(pres_time)
ref_out = sim.data[bp].mean(axis=0).reshape(output_shape.shape())
# Currently, default TensorFlow does not support channels first in conv
use_nengo_dl = nengo_dl is not None and channels_last
ndl_out = np.zeros_like(ref_out)
if use_nengo_dl:
with nengo_dl.Simulator(model, dt=dt) as sim:
sim.run(pres_time)
ndl_out = sim.data[bp].mean(axis=0).reshape(output_shape.shape())
with nengo_loihi.Simulator(model, dt=dt, target='simreal') as sim:
sim.run(pres_time)
real_out = sim.data[bp].mean(axis=0).reshape(output_shape.shape())
with Simulator(model, dt=dt) as sim:
sim.run(pres_time)
sim_out = sim.data[bp].mean(axis=0).reshape(output_shape.shape())
if not output_shape.channels_last:
ref_out = np.transpose(ref_out, (1, 2, 0))
ndl_out = np.transpose(ndl_out, (1, 2, 0))
real_out = np.transpose(real_out, (1, 2, 0))
sim_out = np.transpose(sim_out, (1, 2, 0))
out_max = max(ref_out.max(), sim_out.max())
# --- plot results
rows = 2
cols = 3
ax = plt.subplot(rows, cols, 1)
imshow(test_x, vmin=0, vmax=1, ax=ax)
ax = plt.subplot(rows, cols, 2)
tile(np.transpose(filters[0], (2, 0, 1)), cols=8, ax=ax)
ax = plt.subplot(rows, cols, 3)
plt.hist(ref_out.ravel(), bins=31)
plt.hist(sim_out.ravel(), bins=31)
ax = plt.subplot(rows, cols, 4)
tile(np.transpose(ref_out, (2, 0, 1)), vmin=0, vmax=out_max, cols=8, ax=ax)
ax = plt.subplot(rows, cols, 5)
tile(np.transpose(ndl_out, (2, 0, 1)), vmin=0, vmax=out_max, cols=8, ax=ax)
ax = plt.subplot(rows, cols, 6)
tile(np.transpose(sim_out, (2, 0, 1)), vmin=0, vmax=out_max, cols=8, ax=ax)
if use_nengo_dl:
assert allclose(ndl_out, ref_out, atol=1e-5, rtol=1e-5)
assert allclose(real_out, ref_out, atol=1, rtol=1e-3)
assert allclose(sim_out, ref_out, atol=10, rtol=1e-3)
def test_chip_population_axons(on_chip, precompute, pop_type, channels_last,
Simulator, rng):
"""Check that all types of population axons work as inputs or between cores.
Also, on the chip, dummy axons were still having an effect. Check this is fixed.
"""
def conv_layer(input=None, label=None, **kwargs):
conv = nengo.Convolution(**kwargs)
layer = nengo.Ensemble(conv.output_shape.size, 1, label=label)
conn = (nengo.Connection(input, layer.neurons, transform=conv)
if input is not None else None)
return layer, conv, conn
if pop_type == 16 and not channels_last:
pytest.skip(
"pop16 axons not compatible with single-compartment shifts")
max_rate = 100
amp = 1 / max_rate
n_filters0 = 4
n_filters1 = 4
# 6 x 6 input will have one unused pixel at edge with 3 x 3 kernel and stride 2
input_shape = (6, 6, 1) if channels_last else (1, 6, 6)
input_shape = nengo_transforms.ChannelShape(input_shape,
channels_last=channels_last)
X = rng.uniform(0.2, 1, size=input_shape.shape)
kernel0 = rng.uniform(0.2, 1, size=(1, 1, 1, n_filters0))
kernel1 = rng.uniform(0.1, 0.5, size=(3, 3, n_filters0, n_filters1))
with nengo.Network(seed=0) as net:
nengo_loihi.add_params(net)
net.config[nengo.Ensemble].neuron_type = nengo.SpikingRectifiedLinear(
amplitude=amp)
net.config[nengo.Ensemble].max_rates = nengo.dists.Choice([max_rate])
net.config[nengo.Ensemble].intercepts = nengo.dists.Choice([0])
net.config[nengo.Connection].synapse = 0.005
inp = nengo.Node(X.ravel()) if not on_chip else None
# first layer is off-chip to translate the inputs into spikes
layer0, conv0, _ = conv_layer(
input=inp,
n_filters=n_filters0,
input_shape=input_shape,
channels_last=channels_last,
kernel_size=(1, 1),
init=kernel0,
label="layer0",
)
net.config[layer0].on_chip = on_chip
if on_chip:
assert kernel0.shape[:2] == (1, 1)
w = kernel0[0, 0]
Y = X.dot(w) if channels_last else np.tensordot(w.T, X, axes=1)
layer0.gain = nengo.dists.Choice([0.0])
layer0.bias = Y.ravel() * max_rate
layer1, conv1, conn1 = conv_layer(
input=layer0.neurons,
n_filters=n_filters1,
input_shape=conv0.output_shape,
channels_last=channels_last,
kernel_size=(3, 3),
strides=(2, 2),
init=kernel1,
label="layer1",
)
net.config[conn1].pop_type = pop_type
probe = nengo.Probe(layer1.neurons)
sim_time = 0.1
with Simulator(net, target="sim") as emulator:
emulator.run(sim_time)
with Simulator(net, target="loihi", precompute=precompute) as loihi:
loihi.run(sim_time)
assert np.all(emulator.data[probe].sum(axis=0) > 0)
assert np.array_equal(loihi.data[probe], emulator.data[probe])
train_data_out=train_out[:trainlen]
train_data_out_categorical=to_categorical(train_data_out)
dt = 0.001 # simulation timestep
presentation_time = 0.1 # input presentation time
max_rate = 100 # neuron firing rates
# neuron spike amplitude (scaled so that the overall output is ~1)
amp = 1 / max_rate
# input image shape
input_shape = (1, 7, 7)
n_parallel = 2 # number of parallel network repetitions
minibatch_size = 200
with nengo.Network(seed=0) as net:
# set up the default parameters for ensembles/connections
nengo_loihi.add_params(net)
net.config[nengo.Ensemble].max_rates = nengo.dists.Choice([max_rate])
net.config[nengo.Ensemble].intercepts = nengo.dists.Choice([0])
net.config[nengo.Connection].synapse = None
#neuron_type = nengo.SpikingRectifiedLinear(amplitude=amp)
neuron_type = nengo.LIF(tau_rc=0.02, tau_ref=0.001, amplitude=amp)
#neuron_type = nengo.AdaptiveLIF(amplitude=amp)
#neuron_type = nengo.Izhikevich()
# the input node that will be used to feed in input images
inp = nengo.Node(
nengo.processes.PresentInput(valid_data, presentation_time), size_out=nVariables
)
# the output node provides the 10-dimensional classification
out = nengo.Node(size_in=nClass)
def test_conv_connection(channels, channels_last, Simulator, seed, rng, plt,
allclose):
# load data
with open(os.path.join(test_dir, "mnist10.pkl"), "rb") as f:
test10 = pickle.load(f)
test_x = test10[0][0].reshape((28, 28))
test_x = 1.999 * test_x - 0.999 # range (-1, 1)
input_shape = make_channel_shape(test_x.shape, channels, channels_last)
filters = Gabor(freq=Uniform(0.5, 1)).generate(8, (7, 7), rng=rng)
filters = filters[None, :, :, :] # single channel
filters = np.transpose(filters, (2, 3, 0, 1))
strides = (2, 2)
tau_rc = 0.02
tau_ref = 0.002
tau_s = 0.005
neuron_type = LoihiLIF(tau_rc=tau_rc, tau_ref=tau_ref)
pres_time = 0.1
with nengo.Network(seed=seed) as model:
nengo_loihi.add_params(model)
u = nengo.Node(test_x.ravel(), label="u")
a = nengo.Ensemble(
input_shape.size,
1,
neuron_type=LoihiSpikingRectifiedLinear(),
max_rates=nengo.dists.Choice([40 / channels]),
intercepts=nengo.dists.Choice([0]),
label="a",
)
model.config[a].on_chip = False
if channels == 1:
nengo.Connection(u, a.neurons, transform=1, synapse=None)
elif channels == 2:
# encode image into spikes using two channels (on/off)
if input_shape.channels_last:
nengo.Connection(u, a.neurons[0::2], transform=1, synapse=None)
nengo.Connection(u,
a.neurons[1::2],
transform=-1,
synapse=None)
else:
k = input_shape.spatial_shape[0] * input_shape.spatial_shape[1]
nengo.Connection(u, a.neurons[:k], transform=1, synapse=None)
nengo.Connection(u, a.neurons[k:], transform=-1, synapse=None)
filters = np.concatenate([filters, -filters], axis=2)
else:
raise ValueError("Test not configured for more than two channels")
conv2d_transform = nengo_transforms.Convolution(
8,
input_shape,
strides=strides,
kernel_size=(7, 7),
channels_last=channels_last,
init=filters,
)
output_shape = conv2d_transform.output_shape
gain, bias = neuron_type.gain_bias(max_rates=100, intercepts=0)
gain = gain * 0.01 # account for `a` max_rates
b = nengo.Ensemble(
output_shape.size,
1,
neuron_type=neuron_type,
gain=nengo.dists.Choice([gain[0]]),
bias=nengo.dists.Choice([bias[0]]),
label="b",
)
nengo.Connection(a.neurons,
b.neurons,
synapse=tau_s,
transform=conv2d_transform)
bp = nengo.Probe(b.neurons)
with nengo.Simulator(model, optimize=False) as sim_nengo:
sim_nengo.run(pres_time)
ref_out = sim_nengo.data[bp].mean(axis=0).reshape(output_shape.shape)
with Simulator(model, target="simreal") as sim_emu:
sim_emu.run(pres_time)
emu_out = sim_emu.data[bp].mean(axis=0).reshape(output_shape.shape)
with Simulator(model, hardware_options={"snip_max_spikes_per_step":
800}) as sim_loihi:
sim_loihi.run(pres_time)
sim_out = sim_loihi.data[bp].mean(axis=0).reshape(output_shape.shape)
if not output_shape.channels_last:
ref_out = np.transpose(ref_out, (1, 2, 0))
emu_out = np.transpose(emu_out, (1, 2, 0))
sim_out = np.transpose(sim_out, (1, 2, 0))
out_max = max(ref_out.max(), emu_out.max(), sim_out.max())
# --- plot results
rows = 2
cols = 3
ax = plt.subplot(rows, cols, 1)
imshow(test_x, vmin=0, vmax=1, ax=ax)
ax = plt.subplot(rows, cols, 2)
tile(np.transpose(filters[0], (2, 0, 1)), cols=8, ax=ax)
ax = plt.subplot(rows, cols, 3)
plt.hist(ref_out.ravel(), bins=31)
plt.hist(sim_out.ravel(), bins=31)
ax = plt.subplot(rows, cols, 4)
tile(np.transpose(ref_out, (2, 0, 1)), vmin=0, vmax=out_max, cols=8, ax=ax)
ax = plt.subplot(rows, cols, 5)
tile(np.transpose(emu_out, (2, 0, 1)), vmin=0, vmax=out_max, cols=8, ax=ax)
ax = plt.subplot(rows, cols, 6)
tile(np.transpose(sim_out, (2, 0, 1)), vmin=0, vmax=out_max, cols=8, ax=ax)
assert allclose(emu_out, ref_out, atol=10, rtol=1e-3)
assert allclose(sim_out, ref_out, atol=10, rtol=1e-3)
def test_conv_split(Simulator, rng, plt, allclose):
channels_last = False
# load data
with open(os.path.join(test_dir, 'mnist10.pkl'), 'rb') as f:
test10 = pickle.load(f)
input_shape = nengo_transforms.ChannelShape((1, 28, 28),
channels_last=channels_last)
n_filters = 8
kernel_size = (7, 7)
kernel = Gabor(freq=Uniform(0.5, 1)).generate(n_filters,
kernel_size,
rng=rng)
kernel = kernel[None, :, :, :] # single channel
kernel = np.transpose(kernel, (2, 3, 0, 1))
strides = (2, 2)
seed = 3 # fix seed to do the same computation for both channel positions
with nengo.Network(seed=seed) as net:
nengo_loihi.add_params(net)
a = nengo.Node(test10[0][0].ravel())
# --- make population to turn image into spikes
nc = 1
in_kernel = np.array([1.]).reshape((1, 1, 1, nc))
transform = nengo_transforms.Convolution(1,
input_shape,
kernel_size=(1, 1),
init=in_kernel,
channels_last=channels_last)
b = nengo.Ensemble(transform.output_shape.size,
1,
neuron_type=nengo.SpikingRectifiedLinear(),
max_rates=nengo.dists.Choice([50]),
intercepts=nengo.dists.Choice([0]))
net.config[b].on_chip = False
nengo.Connection(a, b.neurons, transform=transform)
in_shape = transform.output_shape
transform = nengo_transforms.Convolution(n_filters,
in_shape,
kernel_size=kernel_size,
strides=strides,
init=kernel,
channels_last=channels_last)
out_shape = transform.output_shape
split_slices = conv.split_channels(out_shape,
max_size=1024,
max_channels=4)
# --- make convolution population, split across ensembles
cc = []
cp = []
out_shapes = []
xslice = conv.ImageSlice(in_shape)
for yslice in split_slices:
transform_xy = conv.split_transform(transform, xslice, yslice)
out_shapes.append(transform_xy.output_shape)
c = nengo.Ensemble(transform_xy.output_shape.size,
1,
neuron_type=nengo.LIF(),
max_rates=nengo.dists.Choice([15]),
intercepts=nengo.dists.Choice([0]))
nengo.Connection(b.neurons, c.neurons, transform=transform_xy)
cc.append(c)
cp.append(nengo.Probe(c.neurons))
simtime = 0.3
with nengo.Simulator(net, optimize=False) as sim_nengo:
sim_nengo.run(simtime)
hw_opts = dict(snip_max_spikes_per_step=100)
with Simulator(net, seed=seed, hardware_options=hw_opts) as sim_loihi:
sim_loihi.run(simtime)
nengo_out = []
loihi_out = []
for p, out_shape_i in zip(cp, out_shapes):
nengo_out.append(
(sim_nengo.data[p] > 0).sum(axis=0).reshape(out_shape_i.shape))
loihi_out.append(
(sim_loihi.data[p] > 0).sum(axis=0).reshape(out_shape_i.shape))
if channels_last:
nengo_out = np.concatenate(nengo_out, axis=2)
loihi_out = np.concatenate(loihi_out, axis=2)
# put channels first to display them separately
nengo_out = np.transpose(nengo_out, (2, 0, 1))
loihi_out = np.transpose(loihi_out, (2, 0, 1))
else:
nengo_out = np.concatenate(nengo_out, axis=0)
loihi_out = np.concatenate(loihi_out, axis=0)
out_max = np.maximum(nengo_out.max(), loihi_out.max())
# --- plot results
rows = 2
cols = 3
ax = plt.subplot(rows, cols, 1)
imshow(test10[0][0].reshape((28, 28)), vmin=0, vmax=1, ax=ax)
ax = plt.subplot(rows, cols, 2)
tile(np.transpose(kernel[0], (2, 0, 1)), cols=8, ax=ax)
ax = plt.subplot(rows, cols, 3)
plt.hist(nengo_out.ravel(), bins=31)
plt.hist(loihi_out.ravel(), bins=31)
ax = plt.subplot(rows, cols, 4)
tile(nengo_out, vmin=0, vmax=out_max, cols=8, ax=ax)
ax = plt.subplot(rows, cols, 6)
tile(loihi_out, vmin=0, vmax=out_max, cols=8, ax=ax)
assert allclose(loihi_out, nengo_out, atol=0.05 * out_max, rtol=0.15)
def test_conv_connection(channels, channels_last, Simulator, seed, rng, plt,
allclose):
if channels_last:
plt.saveas = None
pytest.xfail("Blocked by CxBase cannot be > 256 bug")
# load data
with open(os.path.join(test_dir, 'mnist10.pkl'), 'rb') as f:
test10 = pickle.load(f)
test_x = test10[0][0].reshape(28, 28)
test_x = 1.999 * test_x - 0.999 # range (-1, 1)
input_shape = nengo_transforms.ChannelShape(
(test_x.shape + (channels, )) if channels_last else
((channels, ) + test_x.shape),
channels_last=channels_last)
filters = Gabor(freq=Uniform(0.5, 1)).generate(8, (7, 7), rng=rng)
filters = filters[None, :, :, :] # single channel
filters = np.transpose(filters, (2, 3, 0, 1))
strides = (2, 2)
tau_rc = 0.02
tau_ref = 0.002
tau_s = 0.005
dt = 0.001
neuron_type = LoihiLIF(tau_rc=tau_rc, tau_ref=tau_ref)
pres_time = 0.1
with nengo.Network(seed=seed) as model:
nengo_loihi.add_params(model)
u = nengo.Node(test_x.ravel(), label="u")
a = nengo.Ensemble(input_shape.size,
1,
neuron_type=LoihiSpikingRectifiedLinear(),
max_rates=nengo.dists.Choice([40 / channels]),
intercepts=nengo.dists.Choice([0]),
label='a')
model.config[a].on_chip = False
if channels == 1:
nengo.Connection(u, a.neurons, transform=1, synapse=None)
elif channels == 2:
# encode image into spikes using two channels (on/off)
if input_shape.channels_last:
nengo.Connection(u, a.neurons[0::2], transform=1, synapse=None)
nengo.Connection(u,
a.neurons[1::2],
transform=-1,
synapse=None)
else:
k = input_shape.spatial_shape[0] * input_shape.spatial_shape[1]
nengo.Connection(u, a.neurons[:k], transform=1, synapse=None)
nengo.Connection(u, a.neurons[k:], transform=-1, synapse=None)
filters = np.concatenate([filters, -filters], axis=2)
else:
raise ValueError("Test not configured for more than two channels")
conv2d_transform = nengo_transforms.Convolution(
8,
input_shape,
strides=strides,
kernel_size=(7, 7),
channels_last=channels_last,
init=filters)
output_shape = conv2d_transform.output_shape
gain, bias = neuron_type.gain_bias(max_rates=100, intercepts=0)
gain = gain * 0.01 # account for `a` max_rates
b = nengo.Ensemble(output_shape.size,
1,
neuron_type=neuron_type,
gain=nengo.dists.Choice([gain[0]]),
bias=nengo.dists.Choice([bias[0]]),
label='b')
nengo.Connection(a.neurons,
b.neurons,
synapse=tau_s,
transform=conv2d_transform)
bp = nengo.Probe(b.neurons)
with nengo.Simulator(model, dt=dt, optimize=False) as sim:
sim.run(pres_time)
ref_out = sim.data[bp].mean(axis=0).reshape(output_shape.shape)
# Currently, non-gpu TensorFlow does not support channels first in conv
use_nengo_dl = HAS_DL and channels_last
ndl_out = np.zeros_like(ref_out)
if use_nengo_dl:
with nengo_dl.Simulator(model, dt=dt) as sim_dl:
sim_dl.run(pres_time)
ndl_out = sim_dl.data[bp].mean(axis=0).reshape(output_shape.shape)
with nengo_loihi.Simulator(model, dt=dt, target='simreal') as sim_real:
sim_real.run(pres_time)
real_out = sim_real.data[bp].mean(axis=0).reshape(output_shape.shape)
with Simulator(model, dt=dt) as sim_loihi:
if "loihi" in sim_loihi.sims:
sim_loihi.sims["loihi"].snip_max_spikes_per_step = 800
sim_loihi.run(pres_time)
sim_out = sim_loihi.data[bp].mean(axis=0).reshape(output_shape.shape)
if not output_shape.channels_last:
ref_out = np.transpose(ref_out, (1, 2, 0))
ndl_out = np.transpose(ndl_out, (1, 2, 0))
real_out = np.transpose(real_out, (1, 2, 0))
sim_out = np.transpose(sim_out, (1, 2, 0))
out_max = max(ref_out.max(), sim_out.max())
# --- plot results
rows = 2
cols = 3
ax = plt.subplot(rows, cols, 1)
imshow(test_x, vmin=0, vmax=1, ax=ax)
ax = plt.subplot(rows, cols, 2)
tile(np.transpose(filters[0], (2, 0, 1)), cols=8, ax=ax)
ax = plt.subplot(rows, cols, 3)
plt.hist(ref_out.ravel(), bins=31)
plt.hist(sim_out.ravel(), bins=31)
ax = plt.subplot(rows, cols, 4)
tile(np.transpose(ref_out, (2, 0, 1)), vmin=0, vmax=out_max, cols=8, ax=ax)
ax = plt.subplot(rows, cols, 5)
tile(np.transpose(ndl_out, (2, 0, 1)), vmin=0, vmax=out_max, cols=8, ax=ax)
ax = plt.subplot(rows, cols, 6)
tile(np.transpose(sim_out, (2, 0, 1)), vmin=0, vmax=out_max, cols=8, ax=ax)
if use_nengo_dl:
assert allclose(ndl_out, ref_out, atol=1e-5, rtol=1e-5)
assert allclose(real_out, ref_out, atol=1, rtol=1e-3)
assert allclose(sim_out, ref_out, atol=10, rtol=1e-3)
def test_split_ensembles(Simulator, seed, rng, plt, allclose):
b_fn = lambda x: x**2
with nengo.Network(seed=seed) as net:
nengo_loihi.add_params(net)
u = nengo.Node(lambda t: np.sin(4 * np.pi * t))
up = nengo.Probe(u, synapse=0.03)
# test auto-splitting large block
a = nengo.Ensemble(1100, 1, label="a")
nengo.Connection(u, a, synapse=None)
ap = nengo.Probe(a, synapse=0.03)
# test connection between two split blocks
b = nengo.Ensemble(400, 1, label="b", seed=seed + 1)
nengo.Connection(a, b, function=b_fn, seed=seed + 2)
net.config[b].block_shape = nengo_loihi.BlockShape((134,), (400,))
bp = nengo.Probe(b, synapse=0.03)
bp10 = nengo.Probe(b[:10], synapse=0.03)
# have one block not be split
c = nengo.Ensemble(400, 1, label="c", seed=seed + 1)
nengo.Connection(a, c, function=b_fn, seed=seed + 2)
cp = nengo.Probe(c, synapse=0.03)
# TODO: uncomment when we allow connections to neuron slices
# ensemble with input to not all blocks, to check synapse splitting,
# specifically the `if len(axon_ids) == 0: continue` in `split_syanpse`.
# However we currently don't allow connections to neuron slices, so ignore.
# d_enc = nengo.dists.UniformHypersphere(surface=True).sample(400, d=1, rng=rng)
# d = nengo.Ensemble(400, 1, label="d", encoders=d_enc)
# net.config[d].block_shape = nengo_loihi.BlockShape((134,), (400,))
# nengo.Connection(a, d.neurons[:200], transform=d_enc[:200])
# nengo.Connection(a, d.neurons[200:], transform=d_enc[200:])
with Simulator(net) as sim:
sim.run(0.5)
assert len(sim.model.objs[a]["out"]) == 2
assert len(sim.model.objs[b]["out"]) == 3
assert len(sim.model.objs[c]["out"]) == 1
y = b_fn(nengo.synapses.Lowpass(0.01).filt(sim.data[up], dt=sim.dt))
plt.plot(sim.trange(), sim.data[up], "k", label="Ideal x")
plt.plot(sim.trange(), sim.data[ap], label="x")
plt.plot(sim.trange(), y, "k", label="Ideal x ** 2")
plt.plot(sim.trange(), sim.data[bp], label="x ** 2, two blocks")
plt.plot(sim.trange(), sim.data[cp], label="x ** 2, one block")
plt.legend()
assert allclose(sim.data[ap], sim.data[up], atol=0.15)
assert allclose(sim.data[bp10], sim.data[bp][:, :10])
# b and c have same seeds, so should be very similar. However, since one
# is split and one is not, discretizing the blocks after splitting means
# that there will be slight numerical differences.
assert allclose(sim.data[cp], sim.data[bp], atol=0.02)
assert allclose(sim.data[bp], y, atol=0.2)
def test_conv_preslice(on_chip, Simulator, plt):
conv2d = pytest.importorskip("nengo._vendor.npconv2d.conv2d")
kernel = np.array([[-1, 2, -1], [-1, 2, -1], [-1, 2, -1]], dtype=float)
kernel /= kernel.max()
image = np.array(
[
[1, 2, 1, 2, 0],
[2, 3, 2, 1, 1],
[1, 2, 1, 2, 3],
[2, 3, 2, 1, 1],
[1, 2, 1, 2, 0],
],
dtype=float,
)
image /= image.max()
image2 = np.column_stack([c * x for c in image.T for x in (1, -1)])
input_gain = 149.0
neuron_type = nengo.SpikingRectifiedLinear()
loihi_neuron = LoihiSpikingRectifiedLinear()
layer0_neuron = loihi_neuron if on_chip else neuron_type
y_ref = layer0_neuron.rates(image.ravel(), input_gain, 0)
y_ref = conv2d.conv2d(y_ref.reshape((1, 5, 5, 1)),
kernel.reshape((3, 3, 1, 1)),
pad="VALID")
y_ref = loihi_neuron.rates(y_ref.ravel(), 1.0, 0.0).reshape((3, 3))
with nengo.Network() as net:
nengo_loihi.add_params(net)
u = nengo.Node(image2.ravel())
a = nengo.Ensemble(
50,
1,
neuron_type=neuron_type,
gain=nengo.dists.Choice([input_gain]),
bias=nengo.dists.Choice([0]),
)
net.config[a].on_chip = on_chip
transform = nengo_transforms.Convolution(n_filters=1,
input_shape=(5, 5, 1),
init=kernel.reshape(
(3, 3, 1, 1)))
b = nengo.Ensemble(
transform.output_shape.size,
1,
neuron_type=neuron_type,
gain=nengo.dists.Choice([1]),
bias=nengo.dists.Choice([0]),
)
nengo.Connection(u, a.neurons, synapse=None)
nengo.Connection(a.neurons[::2], b.neurons, transform=transform)
bp = nengo.Probe(b.neurons, synapse=nengo.Alpha(0.02))
with Simulator(net) as sim:
assert sim.precompute is True
sim.run(0.3)
y_ref = y_ref / input_gain
y = sim.data[bp][-1].reshape((3, -1)) / input_gain
plt.subplot(121)
plt.imshow(y_ref)
plt.colorbar()
plt.subplot(122)
plt.imshow(y)
plt.colorbar()
assert np.allclose(y, y_ref, atol=0.02, rtol=0.1)
def test_seeds(precompute, Simulator, seed):
with nengo.Network(seed=seed) as net:
nengo_loihi.add_params(net)
e0 = nengo.Ensemble(1, 1)
e1 = nengo.Ensemble(1, 1, seed=2)
e2 = nengo.Ensemble(1, 1)
net.config[e2].on_chip = False
nengo.Connection(e0, e1)
nengo.Connection(e0, e2)
with nengo.Network():
n = nengo.Node(0)
e = nengo.Ensemble(1, 1)
nengo.Node(1)
nengo.Connection(n, e)
nengo.Probe(e)
with nengo.Network(seed=8):
nengo.Ensemble(8, 1, seed=3)
nengo.Node(1)
# --- test that seeds are the same as nengo ref simulator
ref = nengo.Simulator(net)
with Simulator(net, precompute=precompute) as sim:
for obj in net.all_objects:
on_chip = (not isinstance(obj, nengo.Node)
and (not isinstance(obj, nengo.Ensemble)
or net.config[obj].on_chip))
seed = sim.model.seeds.get(obj, None)
assert seed is None or seed == ref.model.seeds[obj]
if on_chip:
assert seed is not None
if obj in sim.model.seeded:
assert sim.model.seeded[obj] == ref.model.seeded[obj]
if precompute:
seed0 = sim.sims["host_pre"].model.seeds.get(obj, None)
assert seed0 is None or seed0 == ref.model.seeds[obj]
seed1 = sim.sims["host"].model.seeds.get(obj, None)
assert seed1 is None or seed1 == ref.model.seeds[obj]
else:
seed0 = sim.sims["host"].model.seeds.get(obj, None)
assert seed0 is None or seed0 == ref.model.seeds[obj]
seed1 = None
if not on_chip:
assert seed0 is not None or seed1 is not None
# --- test that seeds that we set are preserved after splitting
model = nengo_loihi.builder.Model()
for i, o in enumerate(net.all_objects):
model.seeds[o] = i
with Simulator(net, model=model, precompute=precompute) as sim:
for i, o in enumerate(net.all_objects):
for name, subsim in sim.sims.items():
if name.startswith("host"):
assert subsim.model.seeds[o] == i