Add firing rate scaling option to Converter

CHANGES.rst

@@ -33,6 +33,7 @@ Release history
 - Added a new ``remove_reset_incs`` graph simplification step. (`#129`_)
 - Added support for UpSampling layers to ``nengo_dl.Converter``. (`#130`_)
 - Added tolerance parameters to ``nengo_dl.Converter.verify``. (`#130`_)
+- Added ``scale_firing_rates`` option to ``nengo_dl.Converter``. (`#134`_)
 
 **Changed**
 
@@ -74,6 +75,7 @@ Release history
 .. _#126: https://github.com/nengo/nengo-dl/pull/126
 .. _#129: https://github.com/nengo/nengo-dl/pull/129
 .. _#130: https://github.com/nengo/nengo-dl/pull/130
+.. _#134: https://github.com/nengo/nengo-dl/pull/134
 
 3.0.0 (December 17, 2019)
 -------------------------

docs/tips.rst

@@ -98,3 +98,20 @@ When debugging spiking performance issues, here are somethings to think about:
    we use both of these techniques.  Again, however, as with any hyperparameters these
    will likely need to be adjusted depending on the application if we want to
    maximize performance.
+4. **Firing rates**. Non-spiking neurons output continuous values every timestep, so
+   it doesn't make much difference whether they are outputting a value of 1 or 100.
+   However, spiking neurons communicate via discrete events, and the rate of those
+   events is proportional to the continuous output value of the corresponding
+   non-spiking counterpart. So a spiking neuron emitting spikes at
+   1Hz is very different than one emitting spikes at 100Hz. Imagine we're
+   simulating the model for 100 timesteps with a simulation timestep of 0.001s. The
+   1Hz neuron is only expected to spike once every 1000 timesteps, so it may not
+   spike at all in our 100 timestep window, meaning that we really have no information
+   about what value that neuron is outputting. Even if a neuron spiked 1 or 2
+   times, that still doesn't provide much information. The 100Hz neuron, on the other
+   hand, would spike about 10 times in our 100 timestep window, allowing us to
+   estimate its firing rate fairly accurately. In conclusion, it is important to look
+   at the firing rates of neurons in your model, and make sure they are spiking fast
+   enough to provide useful information. If they are not spiking fast enough, consider
+   adjusting Ensemble parameterizations (before or after training) or adding
+   regularization terms during training to encourage higher firing rates.

nengo_dl/converter.py

@@ -1,6 +1,7 @@
 """Tools for automatically converting a Keras model to a Nengo network."""
 
 import collections
+import copy
 import logging
 import warnings
 
@@ -64,6 +65,17 @@ class Converter:
         ``{nengo.RectifiedLinear(): nengo.SpikingRectifiedLinear()}``.  Or it can be
         used to swap activation types that don't have a native Nengo implementation,
         e.g. ``{tf.keras.activatons.elu: tf.keras.activations.relu}``.
+    scale_firing_rates: float or dict
+        Scales the inputs of neurons by ``x``, and the outputs by ``1/x``.
+        The idea is that this parameter can be used to increase the firing rates of
+        spiking neurons (by scaling the input), without affecting the overall output
+        (because the output spikes are being scaled back down). Note that this is only
+        strictly true for neuron types with linear activation functions (e.g. ReLU).
+        Nonlinear neuron types (e.g. LIF) will be skewed by this linear scaling on the
+        input/output. ``scale_firing_rates`` can be
+        specified as a float, which will be applied to all layers in the model, or as a
+        dictionary mapping Keras model layers to a scale factor, allowing different
+        scale factors to be applied to different layers.
 
     Attributes
     ----------
@@ -90,12 +102,14 @@ def __init__(
         max_to_avg_pool=False,
         split_shared_weights=False,
         swap_activations=None,
+        scale_firing_rates=None,
     ):
         self.allow_fallback = allow_fallback
         self.inference_only = inference_only
         self.max_to_avg_pool = max_to_avg_pool
         self.split_shared_weights = split_shared_weights
         self.swap_activations = swap_activations or {}
+        self.scale_firing_rates = scale_firing_rates
         self.layer_map = collections.defaultdict(dict)
         self._layer_converters = {}
 
@@ -442,11 +456,47 @@ def add_nengo_obj(self, node_id, biases=None, activation=None):
                     )
             else:
                 # use ensemble to implement the appropriate neuron type
+
+                # apply firing rate scaling
+                if self.converter.scale_firing_rates is None:
+                    scale_firing_rates = None
+                elif isinstance(self.converter.scale_firing_rates, dict):
+                    # look up specific layer rate
+                    scale_firing_rates = self.converter.scale_firing_rates.get(
+                        self.layer, None
+                    )
+                else:
+                    # constant scale applied to all layers
+                    scale_firing_rates = self.converter.scale_firing_rates
+
+                if scale_firing_rates is None:
+                    gain = 1
+                else:
+                    gain = scale_firing_rates
+                    if biases is not None:
+                        biases *= scale_firing_rates
+
+                    if hasattr(activation, "amplitude"):
+                        # copy activation so that we can change amplitude without
+                        # affecting other instances
+                        activation = copy.copy(activation)
+
+                        # bypass read-only protection
+                        type(activation).amplitude.data[
+                            activation
+                        ] /= scale_firing_rates
+                    else:
+                        warnings.warn(
+                            "Firing rate scaling being applied to activation type "
+                            "that does not support amplitude (%s); this will change "
+                            "the output" % type(activation)
+                        )
+
                 obj = nengo.Ensemble(
                     np.prod(self.output_shape(node_id)),
                     1,
                     neuron_type=activation,
-                    gain=nengo.dists.Choice([1]),
+                    gain=nengo.dists.Choice([gain]),
                     bias=nengo.dists.Choice([0]) if biases is None else biases,
                     label=name,
                 ).neurons

nengo_dl/tests/test_converter.py

@@ -570,3 +570,86 @@ def test_upsampling(Layer, size, data_format, rng):
     _test_convert(
         inp, x, inp_vals=[np.arange(np.prod(inp.get_shape())).reshape(inp.shape)],
     )
+
+
+def test_scale_firing_rates():
+    inp = tf.keras.Input(shape=(1,))
+    x = tf.keras.layers.ReLU()(inp)
+    model = tf.keras.Model(inp, x)
+
+    # scaling doesn't affect output at all for non-spiking neurons
+    conv = converter.Converter(model, scale_firing_rates=5)
+    assert conv.verify()
+
+    # works with existing amplitude values
+    neuron_type = nengo.RectifiedLinear(amplitude=2)
+    conv = converter.Converter(
+        model,
+        scale_firing_rates=5,
+        swap_activations={nengo.RectifiedLinear(): neuron_type},
+    )
+    assert neuron_type.amplitude == 2
+    assert conv.net.ensembles[0].neuron_type.amplitude == 2 / 5
+
+    # warning when applying scaling to non-amplitude neuron type
+    inp = tf.keras.Input(shape=(1,))
+    x = tf.keras.layers.Activation(tf.nn.sigmoid)(inp)
+    model = tf.keras.Model(inp, x)
+
+    with pytest.warns(UserWarning, match="does not support amplitude"):
+        conv = converter.Converter(model, scale_firing_rates=5)
+
+    with pytest.raises(ValueError, match="does not match output"):
+        conv.verify()
+
+
+@pytest.mark.parametrize(
+    "scale_firing_rates, expected_rates",
+    [(5, [5, 5]), ({1: 3, 2: 4}, [3, 4]), ({2: 3}, [1, 3])],
+)
+def test_scale_firing_rates_cases(Simulator, scale_firing_rates, expected_rates):
+    input_val = 100
+    bias_val = 50
+    n_steps = 100
+
+    inp = tf.keras.Input(shape=(1,))
+    x0 = tf.keras.layers.ReLU()(inp)
+    x1 = tf.keras.layers.Dense(
+        units=1,
+        activation=tf.nn.relu,
+        kernel_initializer=tf.initializers.constant([[1]]),
+        bias_initializer=tf.initializers.constant([[bias_val]]),
+    )(inp)
+    model = tf.keras.Model(inp, [x0, x1])
+
+    # convert indices to layers
+    scale_firing_rates = (
+        {model.layers[k]: v for k, v in scale_firing_rates.items()}
+        if isinstance(scale_firing_rates, dict)
+        else scale_firing_rates
+    )
+
+    conv = converter.Converter(
+        model,
+        swap_activations={tf.nn.relu: nengo.SpikingRectifiedLinear()},
+        scale_firing_rates=scale_firing_rates,
+    )
+
+    with Simulator(conv.net) as sim:
+        sim.run_steps(
+            n_steps, data={conv.inputs[inp]: np.ones((1, n_steps, 1)) * input_val}
+        )
+
+        for i, p in enumerate(conv.net.probes):
+            # spike heights are scaled down
+            assert np.allclose(np.max(sim.data[p]), 1 / sim.dt / expected_rates[i])
+
+            # number of spikes is scaled up
+            assert np.allclose(
+                np.count_nonzero(sim.data[p]),
+                (input_val if i == 0 else input_val + bias_val)
+                * expected_rates[i]
+                * n_steps
+                * sim.dt,
+                atol=1,
+            )