Getting index error when having an all conv network with larger input size

[kinjalpatel27]

Using an all convolution layer network with nengo-loihi simulator is giving an index error with a larger input size (i.e. any image size > 34x34). Following is the error:

Traceback (most recent call last):
  File "test_loihi.py", line 85, in <module>
    with nengo_loihi.Simulator(
  File "/home/$USER/nengo-loihi/nengo_loihi/simulator.py", line 135, in __init__
    self.model.build(
  File "/home/$USER/nengo-loihi/nengo_loihi/builder/builder.py", line 223, in build
    built = model.builder.build(model, obj, *args, **kwargs)
  File "/home/$USER/nengo/nengo/builder/builder.py", line 239, in build
    return cls.builders[obj_cls](model, obj, *args, **kwargs)
  File "/home/$USER/nengo-loihi/nengo_loihi/builder/network.py", line 44, in build_network
    split_model(model)
  File "/home/$USER/nengo-loihi/nengo_loihi/builder/split_blocks.py", line 130, in split_model
    split_probe(probe, block_map, synapse_map)
  File "/home/$USER/nengo-loihi/nengo_loihi/builder/split_blocks.py", line 199, in split_probe
    new_weights = old_weights[new_weight_inds]
IndexError: index 4 is out of bounds for axis 0 with size 4

With smaller input size (i.e. anything <= 34x34, getting the following error:

  File "test_loihi.py", line 89, in <module>
    sim.run(n_presentations * presentation_time)
  File "/home/$USER/nengo-loihi/nengo_loihi/simulator.py", line 330, in run
    self.run_steps(steps)
  File "/home/$USER/nengo-loihi/nengo_loihi/simulator.py", line 343, in run_steps
    self._runner.run_steps(steps)
  File "/home/$USER/nengo-loihi/nengo_loihi/simulator.py", line 461, in emu_precomputed_host_pre_and_host
    self._chip2host(self.emulator)
  File "/home/$USER/nengo-loihi/nengo_loihi/simulator.py", line 412, in _chip2host
    sim.chip2host(probes_receivers)
  File "/home/$USER/nengo-loihi/nengo_loihi/emulator/interface.py", line 82, in chip2host
    inc = self.probes.send(probe, self._chip2host_sent_steps, receiver)
  File "/home/$USER/nengo-loihi/nengo_loihi/emulator/interface.py", line 638, in send
    x = probe.weight_outputs(outputs)
  File "/home/$USER/nengo-loihi/nengo_loihi/probe.py", line 141, in weight_outputs
    output = output.dot(self.weights[k])
ValueError: shapes (100,1024) and (4,4,3,3) not aligned: 1024 (dim 1) != 3 (dim 2)

Following is the code to reproduce the issue:

import numpy as np
import nengo
from nengo.transforms import ChannelShape
import nengo_loihi

def conv_layer(x, *args, strides=1, activation=True, upsample=False, **kwargs):
    # create a Conv2D transform with the given arguments

    padding = "valid"
    conv = nengo.Convolution(
        *args, strides=(strides, strides), padding=padding, channels_last=True, **kwargs
    )

    if activation:
        # add an ensemble to implement the activation function
        layer = nengo.Ensemble(conv.output_shape.size, 1).neurons
    else:
        # no nonlinearity, so we just use a node
        layer = nengo.Node(size_in=conv.output_shape.size)

    # connect up the input object to the new layer
    nengo.Connection(x, layer, transform=conv)

    # print out the shape information for our new layer
    print("LAYER")
    print(conv.input_shape.shape, "->", conv.output_shape.shape)

    return layer, conv

def model(size=None, channels=3):
    layers = {}
    with nengo.Network(seed=0) as net:
        # the input node that will be used to feed in input images
        nengo_loihi.add_params(net)
        net.config[nengo.Ensemble].neuron_type = nengo.SpikingRectifiedLinear()
        net.config[nengo.Ensemble].max_rates = nengo.dists.Choice([100])
        net.config[nengo.Ensemble].intercepts = nengo.dists.Choice([0])
        net.config[nengo.Connection].synapse = None

        # # the input node that will be used to feed in input images
        # inp = nengo.Node(
        #     nengo.processes.PresentInput(test_images, presentation_time), size_out=28 * 28
        # )
        input_shape = ChannelShape((size, size) + (channels,), channels_last=True)
        input_fn = [0] * input_shape.size
        inp = nengo.Node(input_fn, label="input")
        # the output node provides the 10-dimensional classification
        out = nengo.Node(size_in=10)

        layers["to_spikes"], layer = conv_layer(inp, 4, input_shape, kernel_size=(1, 1))
        x = layers["to_spikes"]
        net.config[layers["to_spikes"].ensemble].on_chip = False

        layers["layer1"], layer1 = conv_layer(
            layers["to_spikes"], 4, layer.output_shape, kernel_size=(3, 3), strides=2,
        )
        x = layers["layer1"]

        layers["layer2"], layer2 = conv_layer(
            x, 4, layer1.output_shape, kernel_size=(3, 3), strides=2, activation=False
        )

    return net, layers, inp

size = 36  # 32

presentation_time = 0.1
n_presentations = 1

net, layers, inp = model(size=size, channels=1)

# Create a random test image
test_images = np.random.normal(size=(1, size * size * 1))

# Setup input
with net:
    inp.output = nengo.processes.PresentInput(
        test_images, presentation_time=presentation_time
    )

# if running on Loihi, increase the max input spikes per step
hw_opts = dict(snip_max_spikes_per_step=120)
with nengo_loihi.Simulator(
    net, dt=0.001, remove_passthrough=False, hardware_options=hw_opts,
) as sim:
    # run the simulation on Loihi
    sim.run(n_presentations * presentation_time)

[hunse]

The crux of the problem here is that you're using a convolutional connection from an ensemble (on-chip) to a node (off-chip). That involves probing the neurons, and then applying the convolutional weights off-chip (and this second part is causing problems).

I'm not sure if this is something that's easy to fix, or even that we want to fix. It might be confusing if you have this convolutional connection that appears to be happening on-chip, but is actually off-chip.

I'll definitely look at getting a better error message at least. But I think you'll want to re-write your network to not have a convolutional connection going from something on-chip to something off-chip. The easiest way to do this would be to make the connection from the ensemble neurons to the node an identity connection (no transform), and then you can do that convolution off-chip by making a connection with the convolution transform from the node do another node.