Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Feat/convnet #1886

Open
wants to merge 8 commits into
base: devel
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
136 changes: 135 additions & 1 deletion psyneulink/core/components/functions/transferfunctions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,\
Expand Down Expand Up @@ -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
# **********************************************************************************************************************
Expand Down
55 changes: 22 additions & 33 deletions psyneulink/core/compositions/composition.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
"""
Expand Down Expand Up @@ -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 [[[]]],
Expand Down
1 change: 1 addition & 0 deletions psyneulink/core/globals/keywords.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
2 changes: 2 additions & 0 deletions psyneulink/core/llvm/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)
Expand Down
75 changes: 75 additions & 0 deletions psyneulink/core/llvm/builtins.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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):
Expand Down
Loading