refactored/revised hebb-deconv syn w/ unit-test

Author: Alexander Ororbia

ngclearn/components/synapses/convolution/convSynapse.py

@@ -50,7 +50,8 @@ 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):
+            batch_size=1, **kwargs
+    ):
         super().__init__(name, **kwargs)
 
         self.filter_init = filter_init

ngclearn/components/synapses/convolution/deconvSynapse.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 deconv2d
+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 deconv2d
 
 class DeconvSynapse(JaxComponent): ## base-level deconvolutional cable
     """
@@ -44,8 +48,10 @@ class DeconvSynapse(JaxComponent): ## base-level deconvolutional 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
@@ -83,29 +89,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
         out = deconv2d(_x, weights, stride_size=stride, padding=padding) * Rscale + biases
         return out
 
-    @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

@@ -92,9 +92,10 @@ def __init__(self, name, shape, x_shape, eta=0., filter_init=None, bias_init=Non
                  stride=1, padding=None, resist_scale=1., w_bound=0.,
                  is_nonnegative=False, w_decay=0., sign_value=1., optim_type="sgd",
                  batch_size=1, **kwargs):
-        super().__init__(name, shape, x_shape=x_shape, filter_init=filter_init,
-                         bias_init=bias_init, resist_scale=resist_scale, stride=stride,
-                         padding=padding, batch_size=batch_size, **kwargs)
+        super().__init__(
+            name, shape, x_shape=x_shape, filter_init=filter_init, bias_init=bias_init, resist_scale=resist_scale,
+            stride=stride, padding=padding, batch_size=batch_size, **kwargs
+        )
 
         self.eta = eta
         self.w_bounds = w_bound

ngclearn/components/synapses/convolution/hebbianDeconvSynapse.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 .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)
@@ -79,13 +86,15 @@ class HebbianDeconvSynapse(DeconvSynapse): ## Hebbian-evolved deconvolutional ca
     """
 
     # Define Functions
-    def __init__(self, name, shape, x_shape, eta=0., filter_init=None, bias_init=None,
-                 stride=1, padding=None, resist_scale=1., w_bound=0., is_nonnegative=False,
-                 w_decay=0., sign_value=1., optim_type="sgd", batch_size=1, **kwargs):
-        super().__init__(name, shape, x_shape=x_shape, filter_init=filter_init,
-                         bias_init=bias_init, resist_scale=resist_scale,
-                         stride=stride, padding=padding, batch_size=batch_size,
-                         **kwargs)
+    def __init__(
+            self, name, shape, x_shape, eta=0., filter_init=None, bias_init=None, stride=1, padding=None,
+            resist_scale=1., w_bound=0., is_nonnegative=False, w_decay=0., sign_value=1., optim_type="sgd",
+            batch_size=1, **kwargs
+    ):
+        super().__init__(
+            name, shape, x_shape=x_shape, filter_init=filter_init, bias_init=bias_init, resist_scale=resist_scale,
+            stride=stride, padding=padding, batch_size=batch_size, **kwargs
+        )
 
         self.eta = eta
         self.w_bounds = w_bound
@@ -112,8 +121,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 = deconv2d(_x, self.weights.value, stride_size=self.stride,
@@ -132,8 +140,7 @@ 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, shape, stride, padding,
-                        delta_shape, pre, post, 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,
@@ -148,10 +155,12 @@ def _compute_update(sign_value, w_decay, bias_init, shape, stride, padding,
             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,
-                shape, stride, padding, delta_shape, pre, post, weights, biases,
-                opt_params):
+    def evolve(
+            opt, sign_value, w_decay, w_bounds, is_nonnegative, bias_init, shape, stride, padding, delta_shape,
+            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
@@ -169,30 +178,19 @@ 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, stride, padding, x_delta_shape, pre, post,
-                      weights):  ## action-backpropagating routine
+    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)
         ## flip sign of back-transmitted signal (if applicable)
         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
@@ -202,14 +200,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_hebbianDeconvSynapse.py

@@ -0,0 +1,110 @@
+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 HebbianDeconvSynapse
+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_HebbianDeconvSynapse1():
+    name = "hebb_deconv_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 = HebbianDeconvSynapse(
+            "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(evolve_process.pure, name="run")
+        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(advance_process.pure, name="run")
+        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(
+        [[[[1.],[2.]],
+          [[2.], [4.]]]]
+    )
+
+    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(
+        [[[[4.]],
+          [[6.]]],
+         [[[6.]],
+          [[9.]]]]
+    )
+    dx_truth = jnp.array(
+        [[[[9.],
+           [6.]],
+          [[6.],
+            [4.]]]]
+    )
+    assert_array_equal(dK, dK_truth)
+    assert_array_equal(dx, dx_truth)
+
+#test_HebbianDeconvSynapse1()