如何在tensorflow中实现stdp？

Question

我正在尝试在tensorflow中实现STDP（峰值时间相关塑性）。有点复杂。有什么想法（完全在张量流图中运行）？在

它是这样工作的：假设我有2个输入神经元，它们通过这个矩阵连接到3个输出神经元：[[1.0, 1.0, 0.0], [0.0, 0.0, 1.0]]（输入神经元0连接到输出神经元0和1…）。在

假设输入神经元有这些峰值（2个神经元，7个时间步）：

Input Spikes:
[[0, 0, 1, 1, 0, 1, 0],
 [1, 1, 0, 0, 0, 0, 1]]

输出神经元的这些峰值（3个神经元，7个时间步）：

^{pr2}$

现在，对于每个非零权重，我要计算一个dw。例如，对于连接到输出神经元0的输入神经元0：

输入神经元0的峰值时间戳是[2, 3, 5]，输出神经元0的时间戳是[3, 6]。现在，我计算所有的增量时间：

Delta Times = [ 2-3, 2-6, 3-3, 3-6, 5-3, 5-6 ] = [ -1, -4, 0, -3, 2, -1 ]

然后，我计算了一些函数（实际的STDP函数，对于这个问题来说并不重要——一些指数函数）

 dw = SUM [ F(-1), F(-4), F(0), F(-3), F(2), F(-1) ]

这是连接输入神经元0和输出神经元0的权重的dw。对所有非零权重重复此操作。在

所以我可以在numpy中完成所有这些，但是我希望能够完全在一个张量流图中完成。尤其是，我一直在计算delta时间。以及如何对所有非零权重并行地执行这些操作。在

这是实际的stdp函数btw（常数可以是参数）：

def stdp_f(x):
    return tf.where(
        x == 0, np.zeros(x.shape), tf.where(
            x > 0, 1.0 * tf.exp(-1.0 * x / 10.0), -1.0 * 1.0 * tf.exp(x / 10.0)))

---

关于性能的一点说明：下面@jdehesa给出的方法既正确又聪明。但结果也很慢。特别是，对于一个由784个输入神经元组成的神经网络，在500多个时间步中，spike_match =步执行（784，1500，1）和（1400，1500）张量的乘法。在

Answer 1

我不熟悉STDP，所以我希望我能正确理解你的意思。我想这就是你所描述的：

import tensorflow as tf

def f(x):
    # STDP function
    return x * 1

def stdp(input_spikes, output_spikes):
    input_shape = tf.shape(input_spikes)
    t = input_shape[-1]
    # Compute STDP function for all possible time difference values
    stdp_values = f(tf.cast(tf.range(-t + 1, t), dtype=input_spikes.dtype))
    # Arrange in matrix such that position [i, j] contains f(i - j)
    matrix_idx = tf.expand_dims(tf.range(t - 1, 2 * t - 1), 1) + tf.range(0, -t, -1)
    stdp_matrix = tf.gather(stdp_values, matrix_idx)
    # Find spike matches
    spike_match = (input_spikes[:, tf.newaxis, :, tf.newaxis] *
                   output_spikes[tf.newaxis, :, tf.newaxis, :])
    # Sum values where there are spike matches
    return tf.reduce_sum(spike_match * stdp_matrix, axis=(2, 3))

# Test
input_spikes = [[0, 0, 1, 1, 0, 1, 0],
                [1, 1, 0, 0, 0, 0, 1]]
output_spikes = [[0, 0, 0, 1, 0, 0, 1],
                 [1, 0, 0, 0, 0, 0, 0],
                 [0, 0, 0, 0, 1, 1, 1]]
with tf.Graph().as_default(), tf.Session() as sess:
    ins = tf.placeholder(tf.float32, [None, None])
    outs = tf.placeholder(tf.float32, [None, None])
    res = stdp(ins, outs)
    res_val = sess.run(res, feed_dict={ins: input_spikes, outs: output_spikes})
    print(res_val)
    # [[ -7.  10. -15.]
    #  [-13.   7. -24.]]

这里我假设f可能很昂贵（而且它的值对于每对神经元都是相同的），所以我只为每个可能的时间增量计算一次，然后在矩阵中重新分配计算的值，这样我就可以在发生输入和输出峰值的坐标对上进行乘法。在

我使用f的identity函数作为占位符，因此结果值实际上只是本例中时间差的总和。在

编辑：仅供参考，将f替换为包含的STDP函数：

^{pr2}$

结果是：

[[-3.4020822   2.1660795  -5.694256  ]
 [-2.974073    0.45364904 -3.1197631 ]]