Add two layer convolutional LCA

src/lava/lib/optimization/solvers/lca/models.py

@@ -9,6 +9,7 @@
 from lava.magma.core.sync.protocols.loihi_protocol import LoihiProtocol
 from lava.magma.core.decorator import implements, tag
 from lava.proc.dense.process import Dense
+from lava.proc.conv.process import Conv
 
 from lava.lib.optimization.solvers.lca.process import LCA1Layer, LCA2Layer
 from lava.lib.optimization.solvers.lca.v1_neuron.process import V1Neuron
@@ -26,18 +27,37 @@ def __init__(self, proc: LCA2Layer):
         weights_exp = proc.weights_exp.get()
         input_val = proc.input.get()
         spike_height = proc.spike_height.get()
-
-        self.v1 = V1Neuron(shape=(weights.shape[0],), tau=T, tau_exp=T_exp,
+        conv = proc.conv.get()
+
+        if conv:
+            weights_T = np.rot90(np.transpose(weights, (3, 1, 2, 0)), 2, (1, 2))
+            kernel_size = weights.shape[1]
+            v1_height = input_val.shape[0] - (kernel_size - 1)
+            v1_width = input_val.shape[1] - (kernel_size - 1)
+            v1_shape = (v1_height, v1_width, weights.shape[0])
+            self.weights_T = Conv(weight=-weights_T, num_message_bits=24,
+                                  weight_exp=weights_exp,
+                                  padding=(kernel_size - 1, kernel_size - 1),
+                                  input_shape=v1_shape)
+            self.weights = Conv(weight=(weights * T), num_message_bits=24,
+                                weight_exp=weights_exp + T_exp,
+                                input_shape=input_val.shape)
+        else:
+            v1_shape = (weights.shape[0],)
+            # weight_exp shifted 8 bits for the weights, 6 for the v1 output.
+            self.weights_T = Dense(weights=-weights.T, num_message_bits=24,
+                                   weight_exp=weights_exp)
+
+            self.weights = Dense(weights=(weights * T), num_message_bits=24,
+                                 weight_exp=weights_exp + T_exp)
+
+
+        self.v1 = V1Neuron(shape=v1_shape, tau=T, tau_exp=T_exp,
                            vth=threshold, two_layer=True)
-        # weight_exp shifted 8 bits for the weights, 6 for the v1 output.
-        self.weights_T = Dense(weights=-weights.T, num_message_bits=24,
-                               weight_exp=weights_exp)
 
-        self.res = ResidualNeuron(shape=(weights.shape[1],),
+        self.res = ResidualNeuron(shape=input_val.shape,
                                   spike_height=spike_height, bias=input_val)
 
-        self.weights = Dense(weights=(weights * T), num_message_bits=24,
-                             weight_exp=weights_exp + T_exp)
 
         self.weights.a_out.connect(self.v1.a_in)
         self.res.s_out.connect(self.weights.s_in)

src/lava/lib/optimization/solvers/lca/process.py

@@ -93,6 +93,7 @@ def __init__(
             tau: ty.Optional[float] = 0.1,
             tau_exp: ty.Optional[int] = 0,
             spike_height: ty.Optional[int] = 1,
+            conv: ty.Optional[bool] = False,
             **kwargs) -> None:
         super().__init__(**kwargs)
 
@@ -101,9 +102,17 @@ def __init__(
         self.tau_exp = Var(shape=(1,), init=tau_exp)
         self.weights = Var(shape=weights.shape, init=weights)
         self.weights_exp = Var(shape=(1,), init=weights_exp)
-        self.v1 = OutPort(shape=(weights.shape[0],))
-        self.res = OutPort(shape=(weights.shape[1],))
-        self.voltage = Var(shape=(weights.shape[0],))
+        if conv:
+            kernel_size = weights.shape[1]
+            v1_height = input_vec.shape[0] - (kernel_size - 1)
+            v1_width = input_vec.shape[1] - (kernel_size - 1)
+            v1_shape = (v1_height, v1_width, weights.shape[0])
+        else:
+            v1_shape = (weights.shape[0],)
+        self.v1 = OutPort(shape=v1_shape)
+        self.res = OutPort(shape=input_vec.shape)
+        self.voltage = Var(shape=v1_shape)
         self.spike_height = Var(shape=(1,), init=spike_height)
         self.input = Var(shape=input_vec.shape, init=input_vec)
         self.input_exp = Var(shape=input_vec.shape, init=input_exp)
+        self.conv = Var(shape=(1,), init=conv)