在 Python 中使用 tensorflow 填充一个权重的张量

Question

我有一个形状为[n, ?, m] 的3d 张量A，沿第三轴有一个非零元素。例如

A[0,0,:] = [0,0,0,1,0,0]
A[1,0,:] = [0,0,1,0,0,0]
A[0,1,:] = [0,1,0,0,0,0]
A[1,1,:] = [1,0,0,0,0,0]

我有一个形状为(1,) 的权重张量w。

我想通过权重 w 膨胀张量 A，以便我可以如下变换张量 A

A[0,0,:] = [0,0,w,1,w,0]
A[1,0,:] = [0,w,1,w,0,0]
A[0,1,:] = [w,1,w,0,0,0]
A[1,1,:] = [1,w,0,0,0,w]

请注意，权重 w 被添加到与非零元素 1 相邻的位置，如果它在边界处，则我们将索引环绕。

如何在 python 中使用 tensorflow 来做到这一点。

Answer 1

编辑：

这是一个更通用的版本，适用于具有多个元素的填充向量：

import tensorflow as tf

def surround_nonzero(a, w):
    # Find non-zero positions
    idx = tf.where(tf.not_equal(a, 0))
    # A vector to shift the last value in the indices
    w_len = tf.shape(w, out_type=tf.int64)[0]
    shift1 = tf.concat([tf.zeros(tf.shape(idx)[-1] - 1, dtype=tf.int64), [1]], axis=0)
    shift_len = shift1 * tf.expand_dims(tf.range(1, w_len + 1), 1)
    # Shift last value of indices using module to wrap around
    a_shape = tf.shape(a, out_type=tf.int64)
    d = a_shape[-1]
    idx_exp = tf.expand_dims(idx, 1)
    idx_prev_exp = (idx_exp - shift_len) % d
    idx_next_exp = (idx_exp + shift_len) % d
    # Reshape shifted indices
    a_rank = tf.rank(a)
    idx_prev = tf.reshape(idx_prev_exp, [-1, a_rank])
    idx_next = tf.reshape(idx_next_exp, [-1, a_rank])
    # Take non-zero values
    nonzero = tf.gather_nd(a, idx)
    # Tile wrapping value twice the number of non-zero values
    n = tf.shape(nonzero)[0]
    w2n = tf.tile(w, [2 * n])
    # Make full index and values for scattering with non-zero values and wrapping value
    idx_full = tf.concat([idx, idx_prev, idx_next], axis=0)
    values_full = tf.concat([nonzero, w2n], axis=0)
    # Make output tensor with scattering
    return tf.scatter_nd(idx_full, values_full, a_shape)

# Test
with tf.Graph().as_default():
    A = tf.constant([[[0, 0, 0, 0, 0, 1, 0, 0],
                      [0, 0, 1, 0, 0, 0, 0, 0]],
                     [[0, 0, 0, 0, 1, 0, 0, 0],
                      [1, 0, 0, 0, 0, 0, 0, 0]]],
                    dtype=tf.int32)
    w = tf.constant([2, 3, 4], dtype=tf.int32)
    out = surround_nonzero(A, w)
    with tf.Session() as sess:
        print(sess.run(out))

输出：

[[[4 0 4 3 2 1 2 3]
  [3 2 1 2 3 4 0 4]]

 [[0 4 3 2 1 2 3 4]
  [1 2 3 4 0 4 3 2]]]

和以前一样，这假设填充始终“适合”，并且不保证填充值重叠的情况下的行为。

---

这是一种使用 tf.scatter_nd 的方法：

import tensorflow as tf

def surround_nonzero(a, w):
    # Find non-zero positions
    idx = tf.where(tf.not_equal(a, 0))
    # A vector to shift the last value in the indices by one
    shift1 = tf.concat([tf.zeros(tf.shape(idx)[-1] - 1, dtype=tf.int64), [1]], axis=0)
    # Shift last value of indices using module to wrap around
    a_shape = tf.shape(a, out_type=tf.int64)
    d = a_shape[-1]
    idx_prev = (idx - shift1) % d
    idx_next = (idx + shift1) % d
    # Take non-zero values
    nonzero = tf.gather_nd(a, idx)
    # Tile wrapping value twice the number of non-zero values
    n = tf.shape(nonzero)[0]
    w2n = tf.tile(w, [2 * n])
    # Make full index and values for scattering with non-zero values and wrapping value
    idx_full = tf.concat([idx, idx_prev, idx_next], axis=0)
    values_full = tf.concat([nonzero, w2n], axis=0)
    # Make output tensor with scattering
    return tf.scatter_nd(idx_full, values_full, a_shape)

# Test
with tf.Graph().as_default():
    A = tf.constant([[[0, 0, 0, 1, 0, 0],
                      [0, 1, 0, 0, 0, 0]],
                     [[0, 0, 1, 0, 0, 0],
                      [1, 0, 0, 0, 0, 0]]],
                    dtype=tf.int32)
    w = tf.constant([2], dtype=tf.int32)
    out = surround_nonzero(A, w)
    with tf.Session() as sess:
        print(sess.run(out))

输出：

[[[0 0 2 1 2 0]
  [2 1 2 0 0 0]]

 [[0 2 1 2 0 0]
  [1 2 0 0 0 2]]]

请注意，这假设每个非零值都被零包围（就像您的情况一样）。否则，scatter 操作会发现重复的索引，并且输出将不是确定性的。