refactored conv/hebb-conv syn w/ unit-test

Author: Alexander Ororbia

ngclearn/components/synapses/convolution/convSynapse.py

@@ -1,10 +1,14 @@
 from jax import random, numpy as jnp, jit
-from ngclearn import resolver, Component, Compartment
 from ngclearn.components.jaxComponent import JaxComponent
-import ngclearn.utils.weight_distribution as dist
-from ngclearn.components.synapses.convolution.ngcconv import conv2d
+from ngcsimlib.compilers.process import transition
+from ngcsimlib.component import Component
+from ngcsimlib.compartment import Compartment
+
+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 conv2d
 
 class ConvSynapse(JaxComponent): ## base-level convolutional cable
     """
@@ -44,8 +48,9 @@ class ConvSynapse(JaxComponent): ## base-level convolutional cable
     """
 
     # Define Functions
-    def __init__(self, name, shape, x_shape, filter_init=None, bias_init=None, stride=1,
-                 padding=None, resist_scale=1., batch_size=1, **kwargs):
+    def __init__(
+            self, name, shape, x_shape, filter_init=None, bias_init=None, stride=1, padding=None, resist_scale=1.,
+            batch_size=1, **kwargs):
         super().__init__(name, **kwargs)
 
         self.filter_init = filter_init
@@ -95,29 +100,22 @@ def __init__(self, name, shape, x_shape, filter_init=None, bias_init=None, strid
                                                          (1, shape[1]))
                                   if bias_init else 0.0)
 
+    @transition(output_compartments=["outputs"])
     @staticmethod
-    def _advance_state(Rscale, padding, stride, weights, biases, inputs):
+    def advance_state(Rscale, padding, stride, weights, biases, inputs):
         _x = inputs
         outputs = conv2d(_x, weights, stride_size=stride, padding=padding) * Rscale + biases
         return outputs
 
-    @resolver(_advance_state)
-    def advance_state(self, outputs):
-        self.outputs.set(outputs)
-
+    @transition(output_compartments=["inputs", "outputs"])
     @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
         outputs = postVals
         return inputs, outputs
 
-    @resolver(_reset)
-    def reset(self, inputs, outputs):
-        self.inputs.set(inputs)
-        self.outputs.set(outputs)
-
     def save(self, directory, **kwargs):
         file_name = directory + "/" + self.name + ".npz"
         if self.bias_init != None:

ngclearn/components/synapses/convolution/hebbianConvSynapse.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,
@@ -143,8 +150,9 @@ def _init(self, batch_size, x_size, shape, stride, padding, pad_args,
         self.x_delta_shape = (dx, dy)
 
     @staticmethod
-    def _compute_update(sign_value, w_decay, bias_init, stride, pad_args,
-                        delta_shape, pre, post, weights):
+    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)
@@ -157,10 +165,12 @@ def _compute_update(sign_value, w_decay, bias_init, stride, pad_args,
             dBiases = jnp.sum(post, axis=0, keepdims=True) * sign_value
         return dWeights, dBiases
 
+    @transition(output_compartments=["opt_params", "weights", "biases", "dWeights", "dBiases"])
     @staticmethod
-    def _evolve(opt, sign_value, w_decay, w_bounds, is_nonnegative, bias_init,
-                stride, pad_args, delta_shape, pre, post, weights, biases,
-                opt_params):
+    def evolve(
+            opt, sign_value, w_decay, w_bounds, is_nonnegative, bias_init, stride, pad_args, delta_shape, pre, post,
+            weights, biases, opt_params
+    ):
         ## calc dFilters / dBiases - update to filters and biases
         dWeights, dBiases = HebbianConvSynapse._compute_update(
             sign_value, w_decay, bias_init, stride, pad_args, delta_shape,
@@ -180,17 +190,11 @@ def _evolve(opt, sign_value, w_decay, w_bounds, is_nonnegative, bias_init,
                 weights = jnp.clip(weights, -w_bounds, w_bounds)
         return opt_params, weights, biases, dWeights, dBiases
 
-    @resolver(_evolve)
-    def evolve(self, opt_params, weights, biases, dWeights, dBiases):
-        self.opt_params.set(opt_params)
-        self.weights.set(weights)
-        self.biases.set(biases)
-        self.dWeights.set(dWeights)
-        self.dBiases.set(dBiases)
-
+    @transition(output_compartments=["dInputs"])
     @staticmethod
-    def _backtransmit(sign_value, x_size, shape, stride, padding, x_delta_shape,
-                      antiPad, post, weights): ## action-backpropagating routine
+    def backtransmit(
+            sign_value, x_size, shape, stride, padding, x_delta_shape, antiPad, post, 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
@@ -206,12 +210,9 @@ def _backtransmit(sign_value, x_size, shape, stride, padding, x_delta_shape,
         dInputs = dInputs * sign_value
         return dInputs
 
-    @resolver(_backtransmit)
-    def backtransmit(self, dInputs):
-        self.dInputs.set(dInputs)
-
+    @transition(output_compartments=["inputs", "outputs", "pre", "post", "dInputs"])
     @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
@@ -221,14 +222,6 @@ def _reset(in_shape, out_shape):
         dInputs = preVals
         return inputs, outputs, pre, post, dInputs
 
-    @resolver(_reset)
-    def reset(self, inputs, outputs, pre, post, dInputs):
-        self.inputs.set(inputs)
-        self.outputs.set(outputs)
-        self.pre.set(pre)
-        self.post.set(post)
-        self.dInputs.set(dInputs)
-
     @classmethod
     def help(cls): ## component help function
         properties = {

tests/components/synapses/convolution/test_hebbianConvSynapse.py

@@ -0,0 +1,108 @@
+from jax import numpy as jnp, random, jit
+from ngcsimlib.context import Context
+import numpy as np
+np.random.seed(42)
+from ngclearn.components import HebbianConvSynapse
+import ngclearn.utils.weight_distribution as dist
+from ngcsimlib.compilers import compile_command, wrap_command
+from numpy.testing import assert_array_equal
+
+from ngcsimlib.compilers.process import Process, transition
+from ngcsimlib.component import Component
+from ngcsimlib.compartment import Compartment
+from ngcsimlib.context import Context
+
+def test_HebbianConvSynapse1():
+    name = "hebb_conv_ctx"
+    ## create seeding keys
+    dkey = random.PRNGKey(1234)
+    dkey, *subkeys = random.split(dkey, 6)
+    dt = 1.  # ms
+    padding_style = "SAME"
+    stride = 1 #2
+    batch_size = 1
+    w_size = 2
+    n_in_chan = 1 #4  # 1
+    n_out_chan = 1 #5  # 1
+    x_size = 2 #4
+
+    shape = (w_size, w_size, n_in_chan, n_out_chan)
+    x_shape = (batch_size, x_size, x_size, n_in_chan)
+
+    # ---- build a simple Hebb-Conv system ----
+    with Context(name) as ctx:
+        a = HebbianConvSynapse(
+            "a", shape, (x_size, x_size), eta=0., filter_init=dist.constant(value=1.), bias_init=None,
+            stride=stride, padding=padding_style, batch_size=batch_size, key=subkeys[0]
+        )
+
+        #"""
+        evolve_process = (Process()
+                          >> a.evolve)
+        ctx.wrap_and_add_command(jit(evolve_process.pure), name="adapt")
+
+        backtransmit_process = (Process()
+                                >> a.backtransmit)
+        ctx.wrap_and_add_command(jit(backtransmit_process.pure), name="backtransmit")
+
+        advance_process = (Process()
+                           >> a.advance_state)
+        ctx.wrap_and_add_command(jit(advance_process.pure), name="run")
+
+        reset_process = (Process()
+                         >> a.reset)
+        ctx.wrap_and_add_command(jit(reset_process.pure), name="reset")
+        #"""
+
+        """
+        reset_cmd, reset_args = ctx.compile_by_key(a, compile_key="reset")
+        ctx.add_command(wrap_command(jit(ctx.reset)), name="reset")
+        advance_cmd, advance_args = ctx.compile_by_key(a, compile_key="advance_state")
+        ctx.add_command(wrap_command(jit(ctx.advance_state)), name="run")
+        evolve_cmd, evolve_args = ctx.compile_by_key(a, compile_key="evolve")
+        ctx.add_command(wrap_command(jit(ctx.evolve)), name="adapt")
+        backpass_cmd, backpass_args = ctx.compile_by_key(a, compile_key="backtransmit")
+        ctx.add_command(wrap_command(jit(ctx.backtransmit)), name="backtransmit")
+        """
+
+    x = jnp.ones(x_shape)
+
+    ctx.reset()
+    a.inputs.set(x)
+    ctx.run(t=1., dt=dt)
+    y = a.outputs.value
+
+    y_truth = jnp.array(
+        [[[[4.],[2.]],
+          [[2.], [1.]]]]
+    )
+
+    assert_array_equal(y, y_truth)
+    # print(y)
+    # print("======")
+
+    # print("NGC-Learn.shape = ", node.outputs.value.shape)
+    a.pre.set(x)
+    a.post.set(y)
+    ctx.adapt(t=1., dt=dt)
+    dK = a.dWeights.value
+    #print(dK)
+    ctx.backtransmit(t=1., dt=dt)
+    dx = a.dInputs.value
+    #print(dx)
+    dK_truth = jnp.array(
+        [[[[9.]],
+          [[6.]]],
+         [[[6.]],
+          [[4.]]]]
+    )
+    dx_truth = jnp.array(
+        [[[[4.],
+           [6.]],
+          [[6.],
+            [9.]]]]
+    )
+    assert_array_equal(dK, dK_truth)
+    assert_array_equal(dx, dx_truth)
+
+#test_HebbianConvSynapse1()