def test_pop_tiny(pop_type, channels_last, nc, request, plt, seed, allclose):
tau_rc = 0.02
tau_ref = 0.001
tau_s = 0.0
dt = 0.001
neuron_bias = 1.
pres_time = 0.4
sti, stj = 1, 1
if nc == 1:
filters = np.array([[-0.5, 2., -0.25], [-0.75, 2., -1.0],
[-0.5, 3., -0.5], [-1.0, 6.,
-0.25]]).reshape(1, 4, 1, 3)
inp_biases = np.array([[1, 5, 1], [2, 1, 2]])
inp_biases = inp_biases[:, :, None]
elif nc == 2:
filters = np.array([[[-0.5, 2., -0.2], [-0.7, 2., -1.0],
[-0.5, 3., -0.5], [-1.0, 6., -0.2]],
[[-1.0, 2., -1.0], [-0.5, 2., -0.5],
[-0.8, 3., -0.2], [-1.0, 4.,
-0.2]]]).reshape(2, 4, 1, 3)
inp_biases = np.array([[[1, 5, 1], [2, 1, 2]], [[0, 3, 1], [4, 2, 1]]])
inp_biases = np.transpose(inp_biases, (1, 2, 0))
# rearrange to (kernel_rows, kernel_cols, in_channels, out_channels)
filters = np.transpose(filters, (2, 3, 0, 1))
inp_biases = inp_biases / (inp_biases.max() + 0.001)
# --- compute nengo_loihi outputs
ni, nj, nk = inp_biases.shape
si, sj, nc, nf = filters.shape
nij = ni * nj
nyi = 1 + (ni - si) // sti
nyj = 1 + (nj - sj) // stj
out_size = nyi * nyj * nf
assert out_size <= 1024
model = Model()
# input block
inp = LoihiBlock(ni * nj * nk, label='inp')
assert inp.n_neurons <= 1024
inp.compartment.configure_relu()
inp.compartment.bias[:] = inp_biases.ravel()
inp_ax = Axon(nij, label='inp_ax')
# we always compute the pixel/channel idxs with channels_last=True
# (not sure why?), and then set it to the correct value afterwards
inp_shape = nengo_transforms.ChannelShape((ni, nj, nk), channels_last=True)
inp_ax.set_compartment_axon_map(target_axons=conv.pixel_idxs(inp_shape),
atoms=conv.channel_idxs(inp_shape))
inp_shape.shape = (ni, nj, nk) if channels_last else (nk, ni, nj)
inp_shape.channels_last = channels_last
inp.add_axon(inp_ax)
model.add_block(inp)
# conv block
neurons = LoihiBlock(out_size, label='neurons')
assert neurons.n_neurons <= 1024
neurons.compartment.configure_lif(tau_rc=tau_rc, tau_ref=tau_ref, dt=dt)
neurons.compartment.configure_filter(tau_s, dt=dt)
neurons.compartment.bias[:] = neuron_bias
synapse = Synapse(np.prod(inp_shape.spatial_shape), label='synapse')
conv2d_transform = nengo_transforms.Convolution(
nf,
inp_shape,
strides=(sti, stj),
channels_last=channels_last,
init=filters,
kernel_size=(1, 3))
weights, indices, axon_to_weight_map, bases = conv.conv2d_loihi_weights(
conv2d_transform)
synapse.set_population_weights(weights,
indices,
axon_to_weight_map,
bases,
pop_type=pop_type)
neurons.add_synapse(synapse)
out_probe = Probe(target=neurons, key='spiked')
neurons.add_probe(out_probe)
inp_ax.target = synapse
model.add_block(neurons)
# simulation
discretize_model(model)
n_steps = int(pres_time / dt)
target = request.config.getoption("--target")
if target == 'loihi':
with HardwareInterface(model, use_snips=False, seed=seed) as sim:
sim.run_steps(n_steps)
sim_out = sim.get_probe_output(out_probe)
else:
with EmulatorInterface(model, seed=seed) as sim:
sim.run_steps(n_steps)
sim_out = sim.get_probe_output(out_probe)
sim_out = np.sum(sim_out, axis=0) * (dt / pres_time)
if channels_last:
sim_out.shape = (nyi, nyj, nf)
sim_out = np.transpose(sim_out, (2, 0, 1))
else:
sim_out.shape = (nf, nyi, nyj)
out_max = sim_out.max()
# --- plot results
rows = 1
cols = 2
ax = plt.subplot(rows, cols, 1)
plt.hist(sim_out.ravel(), bins=11)
ax = plt.subplot(rows, cols, 2)
tile(sim_out, vmin=0, vmax=out_max, grid=True, ax=ax)
# ref_out determined by emulator running code known to work
if nc == 1:
ref_out = np.array([[0.06, 0.02], [0.055, 0.], [0.0825, 0.0225],
[0.125, 0.04]])
elif nc == 2:
ref_out = np.array([[0.0975, 0.02], [0.0825, 0.02], [0.125, 0.055],
[0.2475, 0.0825]])
assert allclose(sim_out[:, :, 0], ref_out, rtol=0, atol=1e-7)
def build_conv2d_connection(model, transform, conn):
assert is_transform_type(transform, ("Convolution", "ConvolutionTranspose"))
if transform.dimensions != 2:
raise NotImplementedError("nengo-loihi only supports 2D convolution")
# Create random number generator
rng = np.random.RandomState(model.seeds[conn])
pre_obj = model.objs[conn.pre_obj]["out"]
post_obj = model.objs[conn.post_obj]["in"]
assert isinstance(pre_obj, (LoihiInput, LoihiBlock))
assert isinstance(post_obj, LoihiBlock)
tau_s = 0.0
if isinstance(conn.synapse, nengo.synapses.Lowpass):
tau_s = conn.synapse.tau
elif conn.synapse is not None:
raise NotImplementedError("Cannot handle non-Lowpass synapses")
# --- pre
assert isinstance(conn.pre_obj, (Neurons, ChipReceiveNeurons))
kernel = transform.sample(rng=rng)
input_shape = transform.input_shape
# Account for nengo spike height of 1/dt
kernel = kernel / model.dt
if isinstance(conn.pre_obj, ChipReceiveNeurons):
neuron_type = conn.pre_obj.neuron_type
elif isinstance(conn.pre_obj, Neurons):
neuron_type = conn.pre_obj.ensemble.neuron_type
if neuron_type is not None and hasattr(neuron_type, "amplitude"):
kernel = kernel * neuron_type.amplitude
# --- post
assert isinstance(conn.post_obj, Neurons)
assert conn.post_slice == slice(None)
gain = model.params[conn.post_obj.ensemble].gain
if not np.all(gain == gain[0]):
# Cannot fold gains into kernel, result would not be convolutional.
# Therefore, Loihi does not support this if we want to share weights.
raise ValidationError(
"All neurons targeted by a Convolution connection must "
"have the same gain",
"gain",
obj=conn.post_obj.ensemble,
)
kernel = kernel * gain[0]
kernel = kernel.astype(nengo.rc.float_dtype)
pop_type = model.config[conn].pop_type
new_transform = copy.copy(transform)
type(new_transform).init.data[new_transform] = kernel
weights, indices, axon_to_weight_map, offsets = conv2d_loihi_weights(new_transform)
synapse = Synapse(np.prod(input_shape.spatial_shape), label="conv2d_weights")
synapse.set_population_weights(
weights, indices, axon_to_weight_map, offsets, pop_type=pop_type
)
post_obj.add_synapse(synapse)
model.objs[conn]["weights"] = synapse
if synapse.atom_bits_extra() > 0:
warnings.warn(
"Using more than 32 'populations' (e.g. convolutional filters) with "
"`pop_type=16` axons has not yet been implemented in NxSDK. This feature "
"is therefore emulator-only."
)
target_axons = -np.ones(pre_obj.n_neurons, dtype=np.int32)
target_axons[conn.pre_slice] = pixel_idxs(input_shape)
atoms = np.zeros(pre_obj.n_neurons, dtype=np.int32)
atoms[conn.pre_slice] = channel_idxs(input_shape)
ax = Axon(np.prod(input_shape.spatial_shape), label="conv2d_weights")
ax.target = synapse
ax.set_compartment_axon_map(target_axons, atoms=atoms)
pre_obj.add_axon(ax)
post_obj.compartment.configure_filter(tau_s, dt=model.dt)
model.params[conn] = BuiltConnection(
eval_points=None, solver_info=None, transform=None, weights=kernel
)
def test_conv2d_weights(channels_last, hw_opts, request, plt, seed, rng,
allclose):
def loihi_rates_n(neuron_type, x, gain, bias, dt):
"""Compute Loihi rates on higher dimensional inputs"""
y = x.reshape(-1, x.shape[-1])
gain = np.asarray(gain)
bias = np.asarray(bias)
if gain.ndim == 0:
gain = gain * np.ones(x.shape[-1])
if bias.ndim == 0:
bias = bias * np.ones(x.shape[-1])
rates = loihi_rates(neuron_type, y, gain, bias, dt)
return rates.reshape(*x.shape)
if channels_last:
plt.saveas = None
pytest.xfail("Blocked by CxBase cannot be > 256 bug")
target = request.config.getoption("--target")
if target != 'loihi' and len(hw_opts) > 0:
pytest.skip("Hardware options only available on hardware")
pop_type = 32
# 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 = test_x[3:24, 3:24]
test_x = 1.999 * test_x - 0.999
filters = Gabor(freq=Uniform(0.5, 1)).generate(8, (7, 7), rng=rng)
sti, stj = 2, 2
tau_rc = 0.02
tau_ref = 0.002
tau_s = 0.005
dt = 0.001
encode_type = nengo.SpikingRectifiedLinear()
encode_gain = 1. / dt
encode_bias = 0.
neuron_type = nengo.LIF(tau_rc=tau_rc, tau_ref=tau_ref)
neuron_gain = 1.
neuron_bias = 1.
pres_time = 0.2
# --- compute ideal outputs
def conv_pm(x, kernel):
y0 = scipy.signal.correlate2d(x[0], kernel, mode='valid')[::sti, ::stj]
y1 = scipy.signal.correlate2d(x[1], kernel, mode='valid')[::sti, ::stj]
return [y0, -y1]
ref_out = np.array([test_x, -test_x])
ref_out = loihi_rates_n(encode_type, ref_out, encode_gain, encode_bias, dt)
ref_out = ref_out / encode_gain
ref_out = np.array([conv_pm(ref_out, kernel) for kernel in filters])
ref_out = ref_out.sum(axis=1) # sum positive and negative parts
ref_out = loihi_rates_n(neuron_type, ref_out, neuron_gain, neuron_bias, dt)
# --- compute nengo_loihi outputs
inp_biases = np.stack([test_x, -test_x], axis=-1 if channels_last else 0)
inp_shape = nengo_transforms.ChannelShape(inp_biases.shape,
channels_last=channels_last)
kernel = np.array([filters, -filters]) # two channels, pos and neg
kernel = np.transpose(kernel, (2, 3, 0, 1))
conv2d_transform = nengo_transforms.Convolution(
8,
inp_shape,
strides=(sti, stj),
channels_last=channels_last,
kernel_size=(7, 7),
init=kernel)
out_size = ref_out.size
nf, nyi, nyj = ref_out.shape
assert out_size <= 1024
model = Model()
# input block
inp = LoihiBlock(inp_shape.size, label='inp')
assert inp.n_neurons <= 1024
inp.compartment.configure_relu()
inp.compartment.bias[:] = inp_biases.ravel()
inp_ax = Axon(np.prod(inp_shape.spatial_shape), label='inp_ax')
inp_ax.set_compartment_axon_map(target_axons=conv.pixel_idxs(inp_shape),
atoms=conv.channel_idxs(inp_shape))
inp.add_axon(inp_ax)
model.add_block(inp)
# conv block
neurons = LoihiBlock(out_size, label='neurons')
assert neurons.n_neurons <= 1024
neurons.compartment.configure_lif(tau_rc=tau_rc, tau_ref=tau_ref, dt=dt)
neurons.compartment.configure_filter(tau_s, dt=dt)
neurons.compartment.bias[:] = neuron_bias
synapse = Synapse(np.prod(inp_shape.spatial_shape), label='synapse')
weights, indices, axon_to_weight_map, bases = conv.conv2d_loihi_weights(
conv2d_transform)
synapse.set_population_weights(weights,
indices,
axon_to_weight_map,
bases,
pop_type=pop_type)
neurons.add_synapse(synapse)
out_probe = Probe(target=neurons, key='spiked')
neurons.add_probe(out_probe)
inp_ax.target = synapse
model.add_block(neurons)
# simulation
discretize_model(model)
n_steps = int(pres_time / dt)
if target == 'loihi':
with HardwareInterface(model, use_snips=False, seed=seed,
**hw_opts) as sim:
sim.run_steps(n_steps)
sim_out = sim.get_probe_output(out_probe)
else:
with EmulatorInterface(model, seed=seed) as sim:
sim.run_steps(n_steps)
sim_out = sim.get_probe_output(out_probe)
sim_out = np.sum(sim_out, axis=0) / pres_time
if channels_last:
sim_out.shape = (nyi, nyj, nf)
sim_out = np.transpose(sim_out, (2, 0, 1))
else:
sim_out.shape = (nf, nyi, nyj)
out_max = max(ref_out.max(), sim_out.max())
# --- plot results
rows = 2
cols = 2
ax = plt.subplot(rows, cols, 1)
tile(filters, cols=8, ax=ax)
ax = plt.subplot(rows, cols, 2)
tile(ref_out, vmin=0, vmax=out_max, 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(sim_out, vmin=0, vmax=1, cols=8, ax=ax)
tile(sim_out, vmin=0, vmax=out_max, cols=8, ax=ax)
assert allclose(sim_out, ref_out, atol=10, rtol=1e-3)
def build_conv2d_connection(model, conn):
if nengo_transforms is None:
# It should not be possible to reach this, because this function is
# only called for a Convolution transform, which can exist only if
# nengo_transforms exists.
raise NotImplementedError("Convolution requires newer Nengo")
if conn.transform.dimensions != 2:
raise NotImplementedError("nengo-loihi only supports 2D convolution")
if conn.transform.padding != "valid":
raise NotImplementedError(
"nengo-loihi only supports convolution with 'valid' padding")
# Create random number generator
rng = np.random.RandomState(model.seeds[conn])
pre_obj = model.objs[conn.pre_obj]['out']
post_obj = model.objs[conn.post_obj]['in']
assert isinstance(pre_obj, (LoihiInput, LoihiBlock))
assert isinstance(post_obj, LoihiBlock)
tau_s = 0.0
if isinstance(conn.synapse, nengo.synapses.Lowpass):
tau_s = conn.synapse.tau
elif conn.synapse is not None:
raise NotImplementedError("Cannot handle non-Lowpass synapses")
# --- pre
assert isinstance(conn.pre_obj, (Neurons, ChipReceiveNeurons))
assert isinstance(conn.transform, nengo_transforms.Convolution)
weights = conn.transform.sample(rng=rng)
input_shape = conn.transform.input_shape
# Account for nengo spike height of 1/dt
weights = weights / model.dt
if isinstance(conn.pre_obj, ChipReceiveNeurons):
neuron_type = conn.pre_obj.neuron_type
elif isinstance(conn.pre_obj, Neurons):
neuron_type = conn.pre_obj.ensemble.neuron_type
if neuron_type is not None and hasattr(neuron_type, 'amplitude'):
weights = weights * neuron_type.amplitude
# --- post
assert isinstance(conn.post_obj, Neurons)
assert conn.post_slice == slice(None)
gain = model.params[conn.post_obj.ensemble].gain
if not np.all(gain == gain[0]):
# Cannot fold gains into weights, result would not be convolutional.
# Therefore, Loihi does not support this if we want to share weights.
raise ValidationError(
"All neurons targeted by a Convolution connection must "
"have the same gain",
"gain",
obj=conn.post_obj.ensemble)
weights = weights * gain[0]
pop_type = 32 # TODO: pick this
new_transform = copy.copy(conn.transform)
type(new_transform).init.data[new_transform] = weights
weights, indices, axon_to_weight_map, offsets = conv2d_loihi_weights(
new_transform)
synapse = Synapse(np.prod(input_shape.spatial_shape),
label="conv2d_weights")
synapse.set_population_weights(weights,
indices,
axon_to_weight_map,
offsets,
pop_type=pop_type)
post_obj.add_synapse(synapse)
model.objs[conn]['weights'] = synapse
target_axons = -np.ones(pre_obj.n_neurons, dtype=int)
target_axons[conn.pre_slice] = pixel_idxs(input_shape)
atoms = np.zeros(pre_obj.n_neurons, dtype=int)
atoms[conn.pre_slice] = channel_idxs(input_shape)
ax = Axon(np.prod(input_shape.spatial_shape), label="conv2d_weights")
ax.target = synapse
ax.set_compartment_axon_map(target_axons, atoms=atoms)
pre_obj.add_axon(ax)
post_obj.compartment.configure_filter(tau_s, dt=model.dt)
model.params[conn] = BuiltConnection(eval_points=None,
solver_info=None,
transform=None,
weights=weights)