revised/refactored hebb/stdp conv/deconv syn w/ unit-tests

Author: Alexander Ororbia

ngclearn/components/synapses/convolution/convSynapse.py

@@ -97,9 +97,9 @@ def __init__(
         if self.bias_init is None:
             info(self.name, "is using default bias value of zero (no bias "
                             "kernel provided)!")
-        self.biases = Compartment(dist.initialize_params(subkeys[2], bias_init,
-                                                         (1, shape[1]))
-                                  if bias_init else 0.0)
+        self.biases = Compartment(
+            dist.initialize_params(subkeys[2], bias_init, (1, shape[1])) if bias_init else 0.0
+        )
 
     @transition(output_compartments=["outputs"])
     @staticmethod

ngclearn/components/synapses/convolution/hebbianConvSynapse.py

@@ -133,8 +133,7 @@ def __init__(self, name, shape, x_shape, eta=0., filter_init=None, bias_init=Non
             [self.weights.value, self.biases.value]
             if bias_init else [self.weights.value]))
 
-    def _init(self, batch_size, x_size, shape, stride, padding, pad_args,
-              weights):
+    def _init(self, batch_size, x_size, shape, stride, padding, pad_args, weights):
         k_size, k_size, n_in_chan, n_out_chan = shape
         _x = jnp.zeros((batch_size, x_size, x_size, n_in_chan))
         _d = conv2d(_x, weights.value, stride_size=stride, padding=padding) * 0
@@ -155,8 +154,7 @@ def _compute_update(
             sign_value, w_decay, bias_init, stride, pad_args, delta_shape, pre, post, weights
     ): ## synaptic kernel adjustment calculation co-routine
         ## compute adjustment to filters
-        dWeights = calc_dK_conv(pre, post, delta_shape=delta_shape,
-                                stride_size=stride, padding=pad_args)
+        dWeights = calc_dK_conv(pre, post, delta_shape=delta_shape, stride_size=stride, padding=pad_args)
         dWeights = dWeights * sign_value
         if w_decay > 0.:  ## apply synaptic decay
             dWeights = dWeights - weights * w_decay
@@ -174,12 +172,10 @@ def evolve(
     ):
         ## calc dFilters / dBiases - update to filters and biases
         dWeights, dBiases = HebbianConvSynapse._compute_update(
-            sign_value, w_decay, bias_init, stride, pad_args, delta_shape,
-            pre, post, weights
+            sign_value, w_decay, bias_init, stride, pad_args, delta_shape, pre, post, weights
         )
         if bias_init != None:
-            opt_params, [weights, biases] = opt(opt_params, [weights, biases],
-                                                [dWeights, dBiases])
+            opt_params, [weights, biases] = opt(opt_params, [weights, biases], [dWeights, dBiases])
         else: ## ignore dBiases since no biases configured
             opt_params, [weights] = opt(opt_params, [weights], [dWeights])
 
@@ -205,8 +201,7 @@ def backtransmit(
         # elif padding == "VALID":
         #     antiPad = _conv_valid_transpose_padding(post.shape[1], x_size,
         #                                             k_size, stride)
-        dInputs = calc_dX_conv(weights, post, delta_shape=x_delta_shape,
-                               stride_size=stride, anti_padding=antiPad)
+        dInputs = calc_dX_conv(weights, post, delta_shape=x_delta_shape, stride_size=stride, anti_padding=antiPad)
         ## flip sign of back-transmitted signal (if applicable)
         dInputs = dInputs * sign_value
         return dInputs

ngclearn/components/synapses/convolution/hebbianDeconvSynapse.py

@@ -143,9 +143,9 @@ def _init(self, batch_size, x_size, shape, stride, padding, pad_args, weights):
     def _compute_update(sign_value, w_decay, bias_init, shape, stride, padding, delta_shape, pre, post, weights):
         k_size, k_size, n_in_chan, n_out_chan = shape
         ## compute adjustment to filters
-        dWeights = calc_dK_deconv(pre, post, delta_shape=delta_shape,
-                                  stride_size=stride, out_size=k_size,
-                                  padding=padding)
+        dWeights = calc_dK_deconv(
+            pre, post, delta_shape=delta_shape, stride_size=stride, out_size=k_size, padding=padding
+        )
         dWeights = dWeights * sign_value
         if w_decay > 0.:  ## apply synaptic decay
             dWeights = dWeights - weights * w_decay
@@ -162,12 +162,10 @@ def evolve(
             pre, post, weights, biases, opt_params
     ):
         dWeights, dBiases = HebbianDeconvSynapse._compute_update(
-            sign_value, w_decay, bias_init, shape, stride, padding, delta_shape,
-            pre, post, weights
+            sign_value, w_decay, bias_init, shape, stride, padding, delta_shape, pre, post, weights
         )
         if bias_init != None:
-            opt_params, [weights, biases] = opt(opt_params, [weights, biases],
-                                                [dWeights, dBiases])
+            opt_params, [weights, biases] = opt(opt_params, [weights, biases], [dWeights, dBiases])
         else: ## ignore dBiases since no biases configured
             opt_params, [weights] = opt(opt_params, [weights], [dWeights])
         ## apply any enforced filter constraints
@@ -182,8 +180,7 @@ def evolve(
     @staticmethod
     def backtransmit(sign_value, stride, padding, x_delta_shape, pre, post, weights):  ## action-backpropagating routine
         ## calc dInputs
-        dInputs = calc_dX_deconv(weights, post, delta_shape=x_delta_shape,
-                                 stride_size=stride, padding=padding)
+        dInputs = calc_dX_deconv(weights, post, delta_shape=x_delta_shape, stride_size=stride, padding=padding)
         ## flip sign of back-transmitted signal (if applicable)
         dInputs = dInputs * sign_value
         return dInputs

ngclearn/components/synapses/convolution/traceSTDPConvSynapse.py

@@ -1,6 +1,13 @@
 from jax import random, numpy as jnp, jit
-from ngclearn import resolver, Component, Compartment
+from ngcsimlib.compilers.process import transition
+from ngcsimlib.component import Component
+from ngcsimlib.compartment import Compartment
+
 from .convSynapse import ConvSynapse
+from ngclearn.utils.weight_distribution import initialize_params
+from ngcsimlib.logger import info
+from ngclearn.utils import tensorstats
+import ngclearn.utils.weight_distribution as dist
 from ngclearn.components.synapses.convolution.ngcconv import (_conv_same_transpose_padding,
                                                               _conv_valid_transpose_padding)
 from ngclearn.components.synapses.convolution.ngcconv import (conv2d, _calc_dX_conv,
@@ -67,12 +74,14 @@ class TraceSTDPConvSynapse(ConvSynapse): ## trace-based STDP convolutional cable
     """
 
     # Define Functions
-    def __init__(self, name, shape, x_shape, A_plus, A_minus, eta=0.,
-                 pretrace_target=0., filter_init=None, stride=1, padding=None,
-                 resist_scale=1., w_bound=0., w_decay=0., batch_size=1, **kwargs):
-        super().__init__(name, shape, x_shape=x_shape, filter_init=filter_init,
-                         bias_init=None, resist_scale=resist_scale, stride=stride,
-                         padding=padding, batch_size=batch_size, **kwargs)
+    def __init__(
+            self, name, shape, x_shape, A_plus, A_minus, eta=0., pretrace_target=0., filter_init=None, stride=1,
+            padding=None, resist_scale=1., w_bound=0., w_decay=0., batch_size=1, **kwargs
+    ):
+        super().__init__(
+            name, shape, x_shape=x_shape, filter_init=filter_init, bias_init=None, resist_scale=resist_scale,
+            stride=stride, padding=padding, batch_size=batch_size, **kwargs
+        )
 
         self.eta = eta
         self.w_bound = w_bound ## soft weight constraint
@@ -107,8 +116,7 @@ def __init__(self, name, shape, x_shape, A_plus, A_minus, eta=0.,
                 self.x_size, k_size, stride)
         ########################################################################
 
-    def _init(self, batch_size, x_size, shape, stride, padding, pad_args,
-              weights):
+    def _init(self, batch_size, x_size, shape, stride, padding, pad_args, weights):
         k_size, k_size, n_in_chan, n_out_chan = shape
         _x = jnp.zeros((batch_size, x_size, x_size, n_in_chan))
         _d = conv2d(_x, weights.value, stride_size=stride, padding=padding) * 0
@@ -126,26 +134,28 @@ def _init(self, batch_size, x_size, shape, stride, padding, pad_args,
         self.x_delta_shape = (dx, dy)
 
     @staticmethod
-    def _compute_update(pretrace_target, Aplus, Aminus, stride, pad_args,
-                        delta_shape, preSpike, preTrace, postSpike, postTrace):
+    def _compute_update(
+            pretrace_target, Aplus, Aminus, stride, pad_args, delta_shape, preSpike, preTrace, postSpike, postTrace
+    ):
         ## Compute long-term potentiation to filters
-        dW_ltp = calc_dK_conv(preTrace - pretrace_target, postSpike * Aplus,
-                              delta_shape=delta_shape, stride_size=stride,
-                              padding=pad_args)
+        dW_ltp = calc_dK_conv(
+            preTrace - pretrace_target, postSpike * Aplus, delta_shape=delta_shape, stride_size=stride, padding=pad_args
+        )
         ## Compute long-term depression to filters
-        dW_ltd = -calc_dK_conv(preSpike, postTrace * Aminus,
-                               delta_shape=delta_shape, stride_size=stride,
-                               padding=pad_args)
+        dW_ltd = -calc_dK_conv(
+            preSpike, postTrace * Aminus, delta_shape=delta_shape, stride_size=stride, padding=pad_args
+        )
         dWeights = (dW_ltp + dW_ltd)
         return dWeights
 
+    @transition(output_compartments=["weights", "dWeights"])
     @staticmethod
-    def _evolve(pretrace_target, Aplus, Aminus, w_decay, w_bound,
-                stride, pad_args, delta_shape, preSpike, preTrace, postSpike,
-                postTrace, weights, eta):
+    def evolve(
+            pretrace_target, Aplus, Aminus, w_decay, w_bound, stride, pad_args, delta_shape, preSpike, preTrace,
+            postSpike, postTrace, weights, eta
+    ):
         dWeights = TraceSTDPConvSynapse._compute_update(
-            pretrace_target, Aplus, Aminus, stride, pad_args, delta_shape,
-            preSpike, preTrace, postSpike, postTrace
+            pretrace_target, Aplus, Aminus, stride, pad_args, delta_shape, preSpike, preTrace, postSpike, postTrace
         )
         if w_decay > 0.:  ## apply synaptic decay
             weights = weights + dWeights * eta - weights * w_decay  ## conduct decayed STDP-ascent
@@ -157,14 +167,11 @@ def _evolve(pretrace_target, Aplus, Aminus, w_decay, w_bound,
             weights = jnp.clip(weights, eps, w_bound - eps)
         return weights, dWeights
 
-    @resolver(_evolve)
-    def evolve(self, weights, dWeights):
-        self.weights.set(weights)
-        self.dWeights.set(dWeights)
-
+    @transition(output_compartments=["dInputs"])
     @staticmethod
-    def _backtransmit(x_size, shape, stride, padding, x_delta_shape, antiPad,
-                      postSpike, weights): ## action-backpropagating routine
+    def backtransmit(
+            x_size, shape, stride, padding, x_delta_shape, antiPad, postSpike, weights
+    ): ## action-backpropagating routine
         ## calc dInputs - adjustment w.r.t. input signal
         k_size, k_size, n_in_chan, n_out_chan = shape
         # antiPad = None
@@ -174,16 +181,12 @@ def _backtransmit(x_size, shape, stride, padding, x_delta_shape, antiPad,
         # elif padding == "VALID":
         #     antiPad = _conv_valid_transpose_padding(postSpike.shape[1], x_size,
         #                                             k_size, stride)
-        dInputs = calc_dX_conv(weights, postSpike, delta_shape=x_delta_shape,
-                               stride_size=stride, anti_padding=antiPad)
+        dInputs = calc_dX_conv(weights, postSpike, delta_shape=x_delta_shape, stride_size=stride, anti_padding=antiPad)
         return dInputs
 
-    @resolver(_backtransmit)
-    def backtransmit(self, dInputs):
-        self.dInputs.set(dInputs)
-
+    @transition(output_compartments=["inputs", "outputs", "preSpike", "postSpike", "preTrace", "postTrace"])
     @staticmethod
-    def _reset(in_shape, out_shape):
+    def reset(in_shape, out_shape):
         preVals = jnp.zeros(in_shape)
         postVals = jnp.zeros(out_shape)
         inputs = preVals
@@ -194,15 +197,6 @@ def _reset(in_shape, out_shape):
         postTrace = postVals
         return inputs, outputs, preSpike, postSpike, preTrace, postTrace
 
-    @resolver(_reset)
-    def reset(self, inputs, outputs, preSpike, postSpike, preTrace, postTrace):
-        self.inputs.set(inputs)
-        self.outputs.set(outputs)
-        self.preSpike.set(preSpike)
-        self.postSpike.set(postSpike)
-        self.preTrace.set(preTrace)
-        self.postTrace.set(postTrace)
-
     @classmethod
     def help(cls): ## component help function
         properties = {

ngclearn/components/synapses/convolution/traceSTDPDeconvSynapse.py

@@ -1,9 +1,17 @@
 from jax import random, numpy as jnp, jit
-from ngclearn import resolver, Component, Compartment
+from ngcsimlib.compilers.process import transition
+from ngcsimlib.component import Component
+from ngcsimlib.compartment import Compartment
+
 from .deconvSynapse import DeconvSynapse
+from ngclearn.utils.weight_distribution import initialize_params
+from ngcsimlib.logger import info
+from ngclearn.utils import tensorstats
+import ngclearn.utils.weight_distribution as dist
 from ngclearn.components.synapses.convolution.ngcconv import (deconv2d, _calc_dX_deconv,
                                                               _calc_dK_deconv, calc_dX_deconv,
                                                               calc_dK_deconv)
+from ngclearn.utils.optim import get_opt_init_fn, get_opt_step_fn
 
 class TraceSTDPDeconvSynapse(DeconvSynapse): ## trace-based STDP deconvolutional cable
     """
@@ -65,13 +73,14 @@ class TraceSTDPDeconvSynapse(DeconvSynapse): ## trace-based STDP deconvolutional
     """
 
     # Define Functions
-    def __init__(self, name, shape, x_shape, A_plus, A_minus, eta=0.,
-                 pretrace_target=0., filter_init=None, stride=1, padding=None,
-                 resist_scale=1., w_bound=0., w_decay=0., batch_size=1,
-                 **kwargs):
-        super().__init__(name, shape, x_shape=x_shape, filter_init=filter_init,
-                         bias_init=None, resist_scale=resist_scale, stride=stride,
-                         padding=padding, batch_size=batch_size, **kwargs)
+    def __init__(
+            self, name, shape, x_shape, A_plus, A_minus, eta=0., pretrace_target=0., filter_init=None, stride=1,
+            padding=None, resist_scale=1., w_bound=0., w_decay=0., batch_size=1, **kwargs
+    ):
+        super().__init__(
+            name, shape, x_shape=x_shape, filter_init=filter_init, bias_init=None, resist_scale=resist_scale,
+            stride=stride, padding=padding, batch_size=batch_size, **kwargs
+        )
 
         self.eta = eta
         self.w_bound = w_bound  ## soft weight constraint
@@ -93,12 +102,10 @@ def __init__(self, name, shape, x_shape, A_plus, A_minus, eta=0.,
 
         ########################################################################
         ## Shape error correction -- do shape correction inference (for local updates)
-        self._init(self.batch_size, self.x_size, self.shape, self.stride,
-                   self.padding, self.pad_args, self.weights)
+        self._init(self.batch_size, self.x_size, self.shape, self.stride, self.padding, self.pad_args, self.weights)
         ########################################################################
 
-    def _init(self, batch_size, x_size, shape, stride, padding, pad_args,
-              weights):
+    def _init(self, batch_size, x_size, shape, stride, padding, pad_args, weights):
         k_size, k_size, n_in_chan, n_out_chan = shape
         _x = jnp.zeros((batch_size, x_size, x_size, n_in_chan))
         _d = deconv2d(_x, self.weights.value, stride_size=self.stride,
@@ -117,8 +124,9 @@ def _init(self, batch_size, x_size, shape, stride, padding, pad_args,
         self.x_delta_shape = (dx, dy)
 
     @staticmethod
-    def _compute_update(pretrace_target, Aplus, Aminus, shape, stride, padding,
-                        delta_shape, preSpike, preTrace, postSpike, postTrace):
+    def _compute_update(
+            pretrace_target, Aplus, Aminus, shape, stride, padding, delta_shape, preSpike, preTrace, postSpike, postTrace
+    ):
         k_size, k_size, n_in_chan, n_out_chan = shape
         ## calc dFilters
         dW_ltp = calc_dK_deconv(preTrace - pretrace_target, postSpike * Aplus,
@@ -130,10 +138,12 @@ def _compute_update(pretrace_target, Aplus, Aminus, shape, stride, padding,
         dWeights = (dW_ltp + dW_ltd)
         return dWeights
 
+    @transition(output_compartments=["weights", "dWeights"])
     @staticmethod
-    def _evolve(pretrace_target, Aplus, Aminus, w_decay, w_bound,
-                shape, stride, padding, delta_shape, preSpike, preTrace, postSpike,
-                postTrace, weights, eta):
+    def evolve(
+            pretrace_target, Aplus, Aminus, w_decay, w_bound, shape, stride, padding, delta_shape, preSpike, preTrace,
+            postSpike, postTrace, weights, eta
+    ):
         dWeights = TraceSTDPDeconvSynapse._compute_update(
             pretrace_target, Aplus, Aminus, shape, stride, padding, delta_shape,
             preSpike, preTrace, postSpike, postTrace
@@ -148,25 +158,17 @@ def _evolve(pretrace_target, Aplus, Aminus, w_decay, w_bound,
             weights = jnp.clip(weights, eps, w_bound - eps)
         return weights, dWeights
 
-    @resolver(_evolve)
-    def evolve(self, weights, dWeights):
-        self.weights.set(weights)
-        self.dWeights.set(dWeights)
-
+    @transition(output_compartments=["dInputs"])
     @staticmethod
-    def _backtransmit(stride, padding, x_delta_shape, preSpike, postSpike,
-                      weights):  ## action-backpropagating routine
+    def backtransmit(stride, padding, x_delta_shape, preSpike, postSpike, weights):  ## action-backpropagating routine
         ## calc dInputs
         dInputs = calc_dX_deconv(weights, postSpike, delta_shape=x_delta_shape,
                                  stride_size=stride, padding=padding)
         return dInputs
 
-    @resolver(_backtransmit)
-    def backtransmit(self, dInputs):
-        self.dInputs.set(dInputs)
-
+    @transition(output_compartments=["inputs", "outputs", "preSpike", "postSpike", "preTrace", "postTrace"])
     @staticmethod
-    def _reset(in_shape, out_shape):
+    def reset(in_shape, out_shape):
         preVals = jnp.zeros(in_shape)
         postVals = jnp.zeros(out_shape)
         inputs = preVals
@@ -177,15 +179,6 @@ def _reset(in_shape, out_shape):
         postTrace = postVals
         return inputs, outputs, preSpike, postSpike, preTrace, postTrace
 
-    @resolver(_reset)
-    def reset(self, inputs, outputs, preSpike, postSpike, preTrace, postTrace):
-        self.inputs.set(inputs)
-        self.outputs.set(outputs)
-        self.preSpike.set(preSpike)
-        self.postSpike.set(postSpike)
-        self.preTrace.set(preTrace)
-        self.postTrace.set(postTrace)
-
     @classmethod
     def help(cls): ## component help function
         properties = {