Deep-R

Author: neworderofjamie

examples/eprop/latency_mnist_deep_r.py …mples/event_prop/latency_mnist_deep_r.pyexamples/eprop/latency_mnist_deep_r.py renamed to examples/event_prop/latency_mnist_deep_r.py examples/eprop/latency_mnist_deep_r.py renamed to examples/event_prop/latency_mnist_deep_r.py

@@ -1,64 +1,65 @@
-import matplotlib.pyplot as plt
-import numpy as np
+import logging
 import matplotlib.pyplot as plt
 import mnist
+import numpy as np
 
 from ml_genn import InputLayer, Layer, SequentialNetwork
 from ml_genn.callbacks import Checkpoint
-from ml_genn.compilers import EPropCompiler, InferenceCompiler
-from ml_genn.connectivity import Dense,FixedProbability
+from ml_genn.compilers import EventPropCompiler, InferenceCompiler
+from ml_genn.connectivity import Dense, FixedProbability
 from ml_genn.initializers import Normal
 from ml_genn.neurons import LeakyIntegrate, LeakyIntegrateFire, SpikeInput
+from ml_genn.optimisers import Adam
 from ml_genn.serialisers import Numpy
+from ml_genn.synapses import Exponential
 
 from time import perf_counter
-from ml_genn.utils.data import (calc_latest_spike_time, calc_max_spikes,
-                                log_latency_encode_data)
-
-from ml_genn.compilers.eprop_compiler import default_params
+from ml_genn.utils.data import linear_latency_encode_data
 
 NUM_INPUT = 784
 NUM_HIDDEN = 128
 NUM_OUTPUT = 10
-BATCH_SIZE = 128
+BATCH_SIZE = 32
 NUM_EPOCHS = 10
+EXAMPLE_TIME = 20.0
+DT = 1.0
 SPARSITY = 0.1
 TRAIN = True
 KERNEL_PROFILING = False
 PLOT_REWIRING = True
 
 mnist.datasets_url = "https://storage.googleapis.com/cvdf-datasets/mnist/"
 labels = mnist.train_labels() if TRAIN else mnist.test_labels()
-spikes = log_latency_encode_data(
+spikes = linear_latency_encode_data(
     mnist.train_images() if TRAIN else mnist.test_images(),
-    20.0, 51)
+    EXAMPLE_TIME - (2.0 * DT), 2.0 * DT)
 
 serialiser = Numpy("latency_mnist_checkpoints")
-network = SequentialNetwork(default_params)
+network = SequentialNetwork()
 with network:
     # Populations
-    input = InputLayer(SpikeInput(max_spikes=BATCH_SIZE * calc_max_spikes(spikes)),
+    input = InputLayer(SpikeInput(max_spikes=BATCH_SIZE * NUM_INPUT),
                                   NUM_INPUT)
-    hidden = Layer(FixedProbability(SPARSITY, Normal(sd=1.0 / np.sqrt(NUM_INPUT))), 
-                   LeakyIntegrateFire(v_thresh=0.61, tau_mem=20.0,
-                                      tau_refrac=5.0),
-                   NUM_HIDDEN)
-    output = Layer(Dense(Normal(sd=1.0 / np.sqrt(NUM_HIDDEN))),
-                   LeakyIntegrate(tau_mem=20.0, readout="sum_var"),
-                   NUM_OUTPUT)
+    hidden = Layer(FixedProbability(SPARSITY, Normal(mean=0.0, sd=0.8)), 
+                   LeakyIntegrateFire(v_thresh=1.0, tau_mem=20.0),
+                   NUM_HIDDEN, Exponential(5.0))
+    output = Layer(Dense(Normal(mean=0.2, sd=0.37)),
+                   LeakyIntegrate(tau_mem=20.0, readout="avg_var"),
+                   NUM_OUTPUT, Exponential(5.0))
 
-max_example_timesteps = int(np.ceil(calc_latest_spike_time(spikes)))
+max_example_timesteps = int(np.ceil(EXAMPLE_TIME / DT))
 if TRAIN:
-    compiler = EPropCompiler(example_timesteps=max_example_timesteps,
-                             losses="sparse_categorical_crossentropy",
-                             optimiser="adam", batch_size=BATCH_SIZE,
-                             deep_r_conns=[hidden], deep_r_l1_strength=1E-8,
-                             deep_r_record_rewirings=({} if not PLOT_REWIRING 
-                                                      else {hidden: "in_hid_rewiring"}),
-                             kernel_profiling=KERNEL_PROFILING)
-    compiled_net = compiler.compile(network)
+    compiler = EventPropCompiler(example_timesteps=max_example_timesteps,
+                                 losses="sparse_categorical_crossentropy",
+                                 batch_size=BATCH_SIZE, dt=DT, deep_r_l1_strength=1E-8,
+                                 kernel_profiling=KERNEL_PROFILING)
+    compiled_net = compiler.compile(network, optimisers={"all_connections": {"weight": Adam(1e-2)}},
+                                    deep_r_conns=[hidden], 
+                                    deep_r_record_rewiring=({} if not PLOT_REWIRING 
+                                                            else {hidden: "in_hid_rewiring"}))
 
     with compiled_net:
+        visualise_examples = [0, 32, 64, 96]
         # Evaluate model on numpy dataset
         start_time = perf_counter()
         callbacks = ["batch_progress_bar", Checkpoint(serialiser)]
@@ -67,22 +68,20 @@
                                                num_epochs=NUM_EPOCHS, shuffle=True,
                                                callbacks=callbacks)
         compiled_net.save_connectivity((NUM_EPOCHS - 1,), serialiser)
-        
         end_time = perf_counter()
         print(f"Accuracy = {100 * metrics[output].result}%")
         print(f"Time = {end_time - start_time}s")
 
         if KERNEL_PROFILING:
             print(f"Neuron update time = {compiled_net.genn_model.neuron_update_time}")
             print(f"Presynaptic update time = {compiled_net.genn_model.presynaptic_update_time}")
-            print(f"Synapse dynamics time = {compiled_net.genn_model.synapse_dynamics_time}")
             print(f"Gradient batch reduce time = {compiled_net.genn_model.get_custom_update_time('GradientBatchReduce')}")
             print(f"Gradient learn time = {compiled_net.genn_model.get_custom_update_time('GradientLearn')}")
             print(f"Reset time = {compiled_net.genn_model.get_custom_update_time('Reset')}")
-            print(f"Softmax1 time = {compiled_net.genn_model.get_custom_update_time('Softmax1')}")
-            print(f"Softmax2 time = {compiled_net.genn_model.get_custom_update_time('Softmax2')}")
-            print(f"Softmax3 time = {compiled_net.genn_model.get_custom_update_time('Softmax3')}")
-
+            print(f"Softmax1 time = {compiled_net.genn_model.get_custom_update_time('BatchSoftmax1')}")
+            print(f"Softmax2 time = {compiled_net.genn_model.get_custom_update_time('BatchSoftmax2')}")
+            print(f"Softmax3 time = {compiled_net.genn_model.get_custom_update_time('BatchSoftmax3')}")
+        
         # Plot rewiring curves
         if PLOT_REWIRING:
             fig, axis = plt.subplots()
@@ -96,7 +95,8 @@
     network.load((NUM_EPOCHS - 1,), serialiser)
 
     compiler = InferenceCompiler(evaluate_timesteps=max_example_timesteps,
-                                 batch_size=BATCH_SIZE)
+                                 reset_in_syn_between_batches=True,
+                                 batch_size=BATCH_SIZE, dt=DT)
     compiled_net = compiler.compile(network)
 
     with compiled_net:

ml_genn/ml_genn/callbacks/custom_update.py

@@ -79,4 +79,4 @@ class CustomUpdateOnTimestepEnd(CustomUpdate):
     def on_timestep_end(self, state, timestep):
         if self._custom_update(state, timestep):
             logger.debug(f"Running custom update {self.name} "
-                         f"at start of timestep {timestep}")
+                         f"at start of timestep {timestep}")

ml_genn/ml_genn/compilers/deep_r.py

@@ -157,18 +157,16 @@ def __init__(self, deep_r_2_ccu, key=None):
         self.num_failed_rewirings = deep_r_2_ccu.extra_global_params["NumFailedRewirings"]
         self.key = key
 
-    def set_params(self, data, **kwargs):
-        # Create empty list to hold recorded data
-        data[self.key] = []
-        self._data = data[self.key]
+    def create_state(self, compiled_network, **kwargs):
+        return []
 
-    def on_batch_end(self, batch, metrics):
+    def on_batch_end(self, state, batch, metrics):
         # Read number of rewirings out of view and add to list
-        self._data.append((self.num_rewirings.view[0],
-                           self.num_failed_rewirings.view[0]))
+        state.append((self.num_rewirings.view[0],
+                      self.num_failed_rewirings.view[0]))
 
-    def get_data(self):
-        return self.key, self._data
+    def get_data(self, state):
+        return self.key, state
 
 
 def add_deep_r(synapse_group, genn_model, compiler, l1_strength, 

ml_genn/ml_genn/compilers/event_prop_compiler.py

@@ -3,7 +3,7 @@
 import sympy
 
 from string import Template
-from typing import Mapping, Union, Tuple
+from typing import Mapping, Sequence, Set, Union, Tuple
 from pygenn import (CustomUpdateVarAccess, SynapseMatrixType,
                     VarAccess, VarAccessMode)
 
@@ -13,7 +13,9 @@
 from .ground_truths import GroundTruth
 from .. import Connection, InputLayer, Layer, Population, Network
 from ..callbacks import (Callback, CustomUpdateOnBatchBegin,
-                         CustomUpdateOnBatchEnd, CustomUpdateOnTimestepEnd)
+                         CustomUpdateOnBatchEnd, CustomUpdateOnEpochBegin,
+                         CustomUpdateOnTimestepBegin, 
+                         CustomUpdateOnTimestepEnd)
 from ..communicators import Communicator
 from ..connection import Connection
 from ..losses import (Loss, MeanSquareError, PerNeuronMeanSquareError,
@@ -38,6 +40,7 @@
 from .deep_r import add_deep_r
 from ..utils.auto_tools import solve_ode
 from ..utils.module import get_object, get_object_mapping
+from ..utils.network import get_underlying_conn
 from ..utils.value import is_value_constant
 
 from .compiler import softmax_1_model, softmax_2_model
@@ -294,8 +297,8 @@ def _add_required_wum_psm_parameters(model: AutoSynapseModel,
                              WeightUpdateModel.has_psm_var_ref)
 
 class CompileState:
-    def __init__(self, network: Network, losses, optimisers, 
-                 supported_matrix_type, backend_name):
+    def __init__(self, network: Network, losses, optimisers, deep_r_conns,
+                 deep_r_record_rewiring, supported_matrix_type, backend_name):
         self.backend_name = backend_name
         self._neuron_reset_vars = []
         self._synapse_reset_vars = []
@@ -307,7 +310,11 @@ def __init__(self, network: Network, losses, optimisers,
         self.update_trial_pops = []
         self.adjoint_limit_pops_vars = []
         self.optimisers = {}
-
+        
+        self.deep_r_conns = set(get_underlying_conn(c) for c in deep_r_conns)
+        self.deep_r_record_rewiring = {get_underlying_conn(c): k
+                                       for c, k in deep_r_record_rewiring.items()}
+        
         # Build list of output populations
         readouts = [p for p in network.populations
                     if p.neuron.readout is not None]
@@ -530,6 +537,7 @@ class EventPropCompiler(Compiler):
                                     spike number, strength for overshoot)
         reg_nu_upper:               Target number of hidden neuron
                                     spikes used for regularisation
+        grad_limit:                 TODO
         max_spikes:                 What is the maximum number of spikes each
                                     neuron (input and hidden) can emit each
                                     trial? This is used to allocate memory 
@@ -541,8 +549,9 @@ class EventPropCompiler(Compiler):
         per_timestep_loss:          Should we use the per-timestep or
                                     per-trial loss functions described above?
         dt:                         Simulation timestep [ms]
-        ttfs_alpha                  TODO
-        softmax_temperature         TODO
+        ttfs_alpha:                 TODO
+        softmax_temperature:        TODO
+        deep_r_l1_strength:         TODO
         batch_size:                 What batch size should be used for
                                     training? In our experience, EventProp works
                                     best with modest batch sizes (32-128)
@@ -565,14 +574,13 @@ class EventPropCompiler(Compiler):
     """
 
     def __init__(self, example_timesteps: int, losses,
-                 reg_lambda: Union[float, Tuple[float, float]] = 0.0, reg_nu_upper: float = 0.0,
-                 grad_limit: float = 100.0,
-                 max_spikes: int = 500, 
-                 strict_buffer_checking: bool = False, 
+                 reg_lambda: Union[float, Tuple[float, float]] = 0.0,
+                 reg_nu_upper: float = 0.0, grad_limit: float = 100.0,
+                 max_spikes: int = 500, strict_buffer_checking: bool = False,
                  per_timestep_loss: bool = False, dt: float = 1.0,
                  ttfs_alpha: float = 0.01, softmax_temperature: float = 1.0,
-                 batch_size: int = 1, rng_seed: int = 0,
-                 kernel_profiling: bool = False,
+                 deep_r_l1_strength: float = 0.01, batch_size: int = 1, 
+                 rng_seed: int = 0, kernel_profiling: bool = False,
                  communicator: Communicator = None,
                  **genn_kwargs):
         supported_matrix_types = [SynapseMatrixType.TOEPLITZ,
@@ -625,6 +633,7 @@ def __init__(self, example_timesteps: int, losses,
         self.per_timestep_loss = per_timestep_loss
         self.ttfs_alpha = ttfs_alpha
         self.softmax_temperature = softmax_temperature
+        self.deep_r_l1_strength = deep_r_l1_strength
 
 
     def pre_compile(self, network: Network, 
@@ -633,14 +642,28 @@ def pre_compile(self, network: Network,
         # to training all weights using the adam optimiser with default params
         optimisers = kwargs.get("optimisers", 
                                 {"all_connections": {"weight": "adam"}})
-
+    
         # Check dictionary has been provided
         if not isinstance(optimisers, Mapping):
-            raise RuntimeError("optimisers should be "
+            raise RuntimeError("'optimisers' should be "
                                "specified as a dictionary")
 
+        # Get sequence of connections to apply Deep-R to and 
+        # check provided value is indeed a sequence
+        deep_r_conns = kwargs.get("deep_r_conns", [])
+        if not isinstance(deep_r_conns, (Set, Sequence)):
+            raise RuntimeError("'deep_r_conns' should be "
+                               "specified as a sequence")
+    
+        # Get dictionary of connections to names to record their deep-r
+        # rewiring stats to and check provided value is indeed a mapping
+        deep_r_record_rewiring = kwargs.get("deep_r_record_rewiring", {})
+        if not isinstance(deep_r_record_rewiring, Mapping):
+            raise RuntimeError("'deep_r_record_rewiring' should be "
+                               "specified as a dictionary")
 
         return CompileState(network, self.losses, optimisers,
+                            deep_r_conns, deep_r_record_rewiring,
                             self.supported_matrix_type, 
                             genn_model.backend_name)
 
@@ -1063,7 +1086,7 @@ def create_compiled_network(self, genn_model, neuron_populations: dict,
                     # If connection is in list of those to use Deep-R on
                     gradient_var_ref = create_wu_var_ref(genn_pop, "weightGradient")
                     weight_var_ref = create_wu_var_ref(genn_pop, "weight")
-                    if c in self.deep_r_conns:
+                    if k in compile_state.deep_r_conns:
                         # Add infrastructure
                         deep_r_2_ccu = add_deep_r(genn_pop, genn_model, self,
                                                   self.deep_r_l1_strength, 
@@ -1072,14 +1095,14 @@ def create_compiled_network(self, genn_model, neuron_populations: dict,
 
                         # If we should record rewirings from
                         # this connection, add to list with key
-                        if c in self.deep_r_record_rewirings:
+                        if k in compile_state.deep_r_record_rewiring:
                             deep_r_record_rewirings_ccus.append(
-                                (deep_r_2_ccu, self.deep_r_record_rewirings[c]))
+                                (deep_r_2_ccu, 
+                                 compile_state.deep_r_record_rewiring[k]))
 
                     # Create weight optimiser custom update
                     cu_weight = self._create_optimiser_custom_update(
-                        f"Weight{i}", weight_var_ref,
-                        create_wu_var_ref(genn_pop, gradient_var_ref), 
+                        f"Weight{i}", weight_var_ref, gradient_var_ref,
                         vars["weight"], genn_model)
 
                     # Add custom update to list of optimisers
@@ -1187,6 +1210,7 @@ def create_compiled_network(self, genn_model, neuron_populations: dict,
         base_validate_callbacks = []
 
         # If Deep-R and L1 regularisation are required, add callback
+        deep_r_required = (len(compile_state.deep_r_conns) > 0)
         if deep_r_required and self.deep_r_l1_strength > 0.0:
             base_train_callbacks.append(
                 CustomUpdateOnBatchEnd("DeepRL1", lambda batch: batch > 0))

ml_genn/ml_genn/utils/model.py

@@ -265,7 +265,7 @@ def __init__(self, model, param_vals={}, var_vals={}, pre_var_vals={},
         self.egp_refs = egp_refs
 
     def process(self):
-        return (super(CustomConnectivityUpdateModel, self).process() 
+        return (super(CustomConnectivityUpdateModel, self)._process() 
                 + (self.pre_var_vals,) + (self.post_var_vals,)
                 + (self.var_refs,) + (self.pre_var_refs,) 
                 + (self.post_var_refs,) + (self.egp_refs,))