Feat/convnet

psyneulink/core/components/functions/transferfunctions.py

@@ -61,7 +61,7 @@
 from psyneulink.core.globals.keywords import \
     ADDITIVE_PARAM, ALL, BIAS, EXPONENTIAL_FUNCTION, \
     GAIN, GAUSSIAN_DISTORT_FUNCTION, GAUSSIAN_FUNCTION, HAS_INITIALIZERS, HOLLOW_MATRIX, \
-    IDENTITY_FUNCTION, IDENTITY_MATRIX, INTERCEPT, LEAK, LINEAR_FUNCTION, LINEAR_MATRIX_FUNCTION, LOGISTIC_FUNCTION, \
+    IDENTITY_FUNCTION, IDENTITY_MATRIX, INTERCEPT, LEAK, LINEAR_FUNCTION, CONV2D_FUNCTION, LINEAR_MATRIX_FUNCTION, LOGISTIC_FUNCTION, \
     TANH_FUNCTION, MATRIX_KEYWORD_NAMES, MATRIX, MATRIX_KEYWORD_VALUES, MAX_INDICATOR, MAX_VAL, MULTIPLICATIVE_PARAM, \
     OFF, OFFSET, ON, PER_ITEM, PROB, PRODUCT, OUTPUT_TYPE, PROB_INDICATOR, \
     RATE, RECEIVER, RELU_FUNCTION, SCALE, SLOPE, SOFTMAX_FUNCTION, STANDARD_DEVIATION, SUM,\
@@ -482,6 +482,140 @@ def _is_identity(self, context=None):
         )
 
 
+# **********************************************************************************************************************
+#                                                    Conv2d
+# **********************************************************************************************************************
+
+class Conv2d(TransferFunction):  # -------------------------------------------------------------------------------------
+    componentName = CONV2D_FUNCTION
+
+    class Parameters(TransferFunction.Parameters):
+        kernel = Parameter(np.array([[0]]), modulable=True)
+        stride = Parameter((1, 1), modulable=False)
+        padding = Parameter((0, 0), modulable=False)
+        dilation = Parameter((1, 1), modulable=False)
+
+    @tc.typecheck
+    def __init__(self,
+                 default_variable=None,
+                 kernel=None,
+                 stride=None,
+                 padding=None,
+                 dilation=None,
+                 params=None,
+                 owner=None,
+                 prefs: tc.optional(is_pref_set) = None):
+
+        super().__init__(
+            default_variable=default_variable,
+            kernel=kernel,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            params=params,
+            owner=owner,
+            prefs=prefs,
+        )
+
+    def _gen_llvm_function_body(self, ctx, builder, params, state, arg_in, arg_out, *, tags:frozenset):
+        # Pretend we have one huge array to work on
+        # TODO: should this be invoked in parts?
+        assert isinstance(arg_in.type.pointee, pnlvm.ir.ArrayType)
+
+        kernel_ptr = pnlvm.helpers.get_param_ptr(builder, self, params, "kernel")
+        stride_ptr = pnlvm.helpers.get_param_ptr(builder, self, params, "stride")
+        padding_ptr = pnlvm.helpers.get_param_ptr(builder, self, params, "padding")
+        dilation_ptr = pnlvm.helpers.get_param_ptr(builder, self, params, "dilation")
+
+        height, width = arg_in.type.pointee.count, arg_in.type.pointee.element.count
+        kernel_height, kernel_width = kernel_ptr.type.pointee.count, kernel_ptr.type.pointee.element.count
+        # result shape should have been determined
+        result_height, result_width = arg_out.type.pointee.count, arg_out.type.pointee.element.count
+
+        # zero arg_out
+        builder.store(arg_out.type.pointee(None), arg_out)
+
+        # dereference padding
+        padding_x = builder.fptoui(builder.load(builder.gep(padding_ptr, [ctx.int32_ty(0), ctx.int32_ty(0)])), ctx.int32_ty)
+        padding_y = builder.fptoui(builder.load(builder.gep(padding_ptr, [ctx.int32_ty(0), ctx.int32_ty(1)])), ctx.int32_ty)
+
+        # dereference stride
+        stride_x = builder.fptoui(builder.load(builder.gep(stride_ptr, [ctx.int32_ty(0), ctx.int32_ty(0)])), ctx.int32_ty)
+        stride_y = builder.fptoui(builder.load(builder.gep(stride_ptr, [ctx.int32_ty(0), ctx.int32_ty(1)])), ctx.int32_ty)
+
+        # dereference dilation
+        dilation_x = builder.fptoui(builder.load(builder.gep(dilation_ptr, [ctx.int32_ty(0), ctx.int32_ty(0)])), ctx.int32_ty)
+        dilation_y = builder.fptoui(builder.load(builder.gep(dilation_ptr, [ctx.int32_ty(0), ctx.int32_ty(1)])), ctx.int32_ty)
+
+        def _get_variable_idx(builder, x, y):
+            val_ptr = builder.alloca(ctx.float_ty)
+            x_padding_top = builder.icmp_unsigned('<', x, padding_x)
+            x_padding_bottom = builder.icmp_unsigned('>=', x, builder.add(ctx.int32_ty(height), padding_x))
+            y_padding_left = builder.icmp_unsigned('<', y, padding_y)
+            y_padding_right = builder.icmp_unsigned('>=', y, builder.add(ctx.int32_ty(width), padding_y))
+            pred = builder.or_(x_padding_top, x_padding_bottom)
+            pred2 = builder.or_(y_padding_left, y_padding_right)
+            pred = builder.or_(pred, pred2)
+            with builder.if_else(pred) as (then, otherwise):
+                with then:
+                    # in padding zone
+                    builder.store(ctx.float_ty(-0.0), val_ptr)
+                with otherwise:
+                    x = builder.sub(x, padding_x)
+                    y = builder.sub(y, padding_y)
+                    var_ptr = builder.gep(arg_in, [ctx.int32_ty(0), x, y])
+                    builder.store(builder.load(var_ptr), val_ptr)
+            return val_ptr
+
+        with pnlvm.helpers.for_loop_zero_inc(builder, ctx.int32_ty(result_height), "result_height") as (builder, h):
+            with pnlvm.helpers.for_loop_zero_inc(builder, ctx.int32_ty(result_width), "result_width") as (builder, w):
+                with pnlvm.helpers.for_loop_zero_inc(builder, ctx.int32_ty(kernel_height), "kernel_height") as (builder, k_h):
+                    with pnlvm.helpers.for_loop_zero_inc(builder, ctx.int32_ty(kernel_width), "kernel_width") as (builder, k_w):
+                        i_x = builder.add(builder.mul(k_h, dilation_x), builder.mul(stride_x, h))
+                        i_y = builder.add(builder.mul(k_w, dilation_y), builder.mul(stride_y, w))
+                        variable_ptr = _get_variable_idx(builder, i_x, i_y)
+                        var = builder.load(variable_ptr)
+                        kernel_ptr = builder.gep(kernel_ptr, [ctx.int32_ty(0), k_h, k_w])
+                        kernel = builder.load(kernel_ptr)
+                        result_ptr = builder.gep(arg_out, [ctx.int32_ty(0), h, w])
+                        res = builder.load(result_ptr)
+
+                        builder.store(builder.fadd(res, builder.fmul(var, kernel)), result_ptr)
+
+        return builder
+
+    def _function(self,
+                 variable=None,
+                 context=None,
+                 params=None,
+                 ):
+        kernel = self._get_current_parameter_value("kernel", context)
+        stride = self._get_current_parameter_value("stride", context)
+        padding = self._get_current_parameter_value("padding", context)
+        dilation = self._get_current_parameter_value("dilation", context)
+
+        height, width = variable.shape
+        kernel_height, kernel_width = kernel.shape
+
+        result_height = int(np.floor(((height + 2 *padding[0] - dilation[0] *(kernel_height - 1) - 1) / stride[0]) + 1))
+        result_width = int(np.floor(((width + 2 *padding[1] - dilation[1] *(kernel_width - 1) - 1) / stride[1]) + 1))
+
+        result = np.zeros((result_height, result_width))
+
+        # apply padding
+        variable = np.pad(variable, [(padding[0], padding[0]), (padding[1], padding[1])])
+
+        for h in range(result_height):
+            for w in range(result_width):
+                # apply kernel op.
+                for k_h in range(0, kernel_height):
+                    i_x = k_h * dilation[0] + stride[0] * h
+                    for k_w in range(0, kernel_width):
+                        i_y = k_w * dilation[1] + stride[1] * w
+                        result[h][w] += kernel[k_h][k_w] * variable[i_x][i_y]
+
+        return self.convert_output_type(result)
+
 # **********************************************************************************************************************
 #                                                    Exponential
 # **********************************************************************************************************************

psyneulink/core/compositions/composition.py

@@ -7577,38 +7577,20 @@ def evaluate(
         else:
             return net_outcome
 
-
-    def _infer_target_nodes(self, targets: dict):
-        """
-        Maps targets onto target mechanisms (as needed by learning)
-
-        Returns
-        ---------
-
-        `dict`:
-            Dict mapping TargetMechanisms -> target values
-        """
-        ret = {}
-        for node, values in targets.items():
-            if (NodeRole.TARGET not in self.get_roles_by_node(node)
-                    and NodeRole.LEARNING not in self.get_roles_by_node(node)):
-                node_efferent_mechanisms = [x.receiver.owner for x in node.efferents if x in self.projections]
-                comparators = [x for x in node_efferent_mechanisms
-                               if (isinstance(x, ComparatorMechanism)
-                                   and NodeRole.LEARNING in self.get_roles_by_node(x))]
-                comparator_afferent_mechanisms = [x.sender.owner for c in comparators for x in c.afferents]
-                target_nodes = [t for t in comparator_afferent_mechanisms
-                                if (NodeRole.TARGET in self.get_roles_by_node(t)
-                                    and NodeRole.LEARNING in self.get_roles_by_node(t))]
-
-                if len(target_nodes) != 1:
-                    # Invalid specification: no valid target nodes or ambiguity in which target node to choose
-                    raise Exception(f"Unable to infer learning target node from output node {node} of {self.name}")
-
-                ret[target_nodes[0]] = values
-            else:
-                ret[node] = values
-        return ret
+    def _infer_target_node(self, node):
+        if (NodeRole.TARGET not in self.get_roles_by_node(node) and NodeRole.LEARNING not in self.get_roles_by_node(node)):
+            node_efferent_mechanisms = [x.receiver.owner for x in node.efferents if x in self.projections]
+            comparators = [x for x in node_efferent_mechanisms
+                            if (isinstance(x, ComparatorMechanism)
+                                and NodeRole.LEARNING in self.get_roles_by_node(x))]
+            comparator_afferent_mechanisms = [x.sender.owner for c in comparators for x in c.afferents]
+            target_nodes = [t for t in comparator_afferent_mechanisms
+                            if (NodeRole.TARGET in self.get_roles_by_node(t)
+                                and NodeRole.LEARNING in self.get_roles_by_node(t))]
+
+            if len(target_nodes) != 1:
+                # Invalid specification: no valid target nodes or ambiguity in which target node to choose
+                raise Exception(f"Unable to infer learning target node from output node {node} of {self.name}")
 
     def _parse_learning_spec(self, inputs, targets):
         """
@@ -7648,7 +7630,14 @@ def _recursive_update(d, u):
             return d
 
         if targets is not None:
-            targets = self._infer_target_nodes(targets)
+            inferred_targets = {}
+            for node, values in targets.items():
+                target_node = self._infer_target_node(node)
+                if target_node is not None:
+                    inferred_targets[target_node] = values
+                else:
+                    inferred_targets[node] = values
+            targets = inferred_targets
             inputs = _recursive_update(inputs, targets)
 
         # 3) Resize inputs to be of the form [[[]]],

psyneulink/core/globals/keywords.py

@@ -549,6 +549,7 @@ def _is_metric(metric):
 # TransferFunctions:
 IDENTITY_FUNCTION = 'Identity Function'
 LINEAR_FUNCTION = "Linear Function"
+CONV2D_FUNCTION = "Conv2d Function"
 LEABRA_FUNCTION = "Leabra Function"
 EXPONENTIAL_FUNCTION = "Exponential Function"
 LOGISTIC_FUNCTION = "Logistic Function"

psyneulink/core/llvm/__init__.py

@@ -140,6 +140,7 @@ def init_builtins():
 
         # Matrix/Vector
         builtins.setup_vxm(ctx)
+        builtins.setup_mxm(ctx)
         builtins.setup_vxm_transposed(ctx)
         builtins.setup_vec_add(ctx)
         builtins.setup_vec_sum(ctx)
@@ -148,6 +149,7 @@ def init_builtins():
         builtins.setup_mat_sub(ctx)
         builtins.setup_vec_hadamard(ctx)
         builtins.setup_mat_hadamard(ctx)
+        builtins.setup_vec_outer_product(ctx)
         builtins.setup_vec_scalar_mult(ctx)
         builtins.setup_mat_scalar_mult(ctx)
         builtins.setup_mat_scalar_add(ctx)

psyneulink/core/llvm/builtins.py

@@ -69,6 +69,53 @@ def setup_vxm(ctx):
 
     builder.ret_void()
 
+def setup_mxm(ctx):
+    # Setup types
+    double_ptr_ty = ctx.float_ty.as_pointer()
+    # Arguments
+    # 1) Matrix ptr (X by Y)
+    # 2) Matrix ptr (Y by Z)
+    # 3) X dimension size
+    # 4) Y dimension size
+    # 5) Z dimension size
+    # 6) Output matrix pointer
+    builder = _setup_builtin_func_builder(ctx, "mxm", (double_ptr_ty, double_ptr_ty, ctx.int32_ty, ctx.int32_ty, ctx.int32_ty, double_ptr_ty))
+    m1, m2, x, y, z, o = builder.function.args
+
+    # zero the output matrix
+    with helpers.for_loop_zero_inc(builder, x, "zero_outer") as (b1, index_i):
+        with helpers.for_loop_zero_inc(b1, z, "zero_inner") as (b2, index_j):
+            matrix_index = b2.mul(index_i, z)
+            matrix_index = b2.add(matrix_index, index_j)
+            matrix_ptr = b2.gep(o, [matrix_index])
+            b2.store(ctx.float_ty(0), matrix_ptr)
+
+    # Multiplication
+    with helpers.for_loop_zero_inc(builder, x, "mxm_outer") as (b1, index_i):
+        with helpers.for_loop_zero_inc(b1, y, "mxm_inner") as (b2, index_j):
+            with helpers.for_loop_zero_inc(b2, z, "mxm_inner_2") as (b3, index_k):
+                # Multiplication and accumulation
+                output_index = builder.mul(index_i, z)
+                output_index = builder.add(output_index, index_k)
+                out_ptr = builder.gep(o, [output_index])
+                out_el = builder.load(out_ptr)
+
+                m1_index = builder.mul(index_i, y)
+                m1_index = builder.add(m1_index, index_j)
+                m1_ptr = builder.gep(m1, [m1_index])
+                m1_el = builder.load(m1_ptr)
+
+                m2_index = builder.mul(index_j, z)
+                m2_index = builder.add(m2_index, index_k)
+                m2_ptr = builder.gep(m2, [m2_index])
+                m2_el = builder.load(m2_ptr)
+
+                new_el = builder.fmul(m1_el, m2_el)
+                new_el = builder.fadd(new_el, out_el)
+
+                builder.store(new_el, out_ptr)
+
+    builder.ret_void()
 
 def setup_vxm_transposed(ctx):
     # Setup types
@@ -327,6 +374,34 @@ def setup_mat_hadamard(ctx):
 
     builder.ret_void()
 
+# outer product of vectors
+def setup_vec_outer_product(ctx):
+    # Setup types
+    double_ptr_ty = ctx.float_ty.as_pointer()
+
+    # builtin vector magnitude func
+    # param1: ptr to vec 1
+    # param2: ptr to vec 2
+    # param3: dim_x of vec 1
+    # param4: dim_y of vec 2
+    # param5: output ptr (should be dim_x by dim_y)
+    builder = _setup_builtin_func_builder(ctx, "vec_outer_product", (double_ptr_ty, double_ptr_ty, ctx.int32_ty, ctx.int32_ty, double_ptr_ty))
+    v1, v2, dim_x, dim_y, o = builder.function.args
+
+    with helpers.for_loop_zero_inc(builder, dim_x, "vec_outer_product_outer") as (b1, x):
+        with helpers.for_loop_zero_inc(b1, dim_y, "vec_outer_product_inner") as (b2, y):
+            matrix_index = b2.mul(x, dim_y)
+            matrix_index = b2.add(matrix_index, y)
+            v1_ptr = b2.gep(v1, [x])
+            v2_ptr = b2.gep(v2, [y])
+            o_ptr = b2.gep(o, [matrix_index])
+
+            v1_val = b2.load(v1_ptr)
+            v2_val = b2.load(v2_ptr)
+            o_val = b2.fmul(v1_val, v2_val)
+            b2.store(o_val, o_ptr)
+
+    builder.ret_void()
 
 # matrix subtraction
 def setup_mat_sub(ctx):