TensorFlow - создание тензора из трафарета

Question

При поиске этой проблемы я не нашел ничего связанного, поэтому, пожалуйста, дайте мне знать, если что-то подобное уже существует.

Так вот в чем проблема. Представьте, что у вас есть шаблон Tensor A , содержащий ненулевые значения столбца B . Теперь я хочу сгенерировать Тензор B , содержащий значения A в правильной позиции. A должно быть dim (n, 3), а B dim (2n + 1, n). Разные измерения возникают из разных сеток конечных элементов. Например, пусть n = 3 скажет, что у нас есть A

A = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]

и мы хотим создать Тензор B из A как

B [0,0] = A [0,0], B [1,0] = A [0,1], B [2,0] = A [0,2], B [2,1] = A [1,0], B [3,1] = A [1,1], B [4,1] = A [1,2], B [4,2] = A [2,0], B [5,2] = A [2,1] и B [6,2] = A [2,2].

Так что для примера B будет

B = [[1, 0, 0], [2, 0, 0], [3, 4, 0], [0, 5, 0], [0, 6, 7], [0, 0, 8], [0, 0, 9]].

Кроме того, это должно быть дифференцируемым, если это вообще возможно. Поскольку я хочу обучить Тензор A некоторым параметрам, но используйте B , чтобы вычислить произведение матрицы на вектор для функции потерь. В этот момент я не буду использовать пакетное обучение.

Я пытался использовать функцию tf.assign () , чтобы назначить каждое значение для Tensor B . Я ищу лучший вариант, так как он становится очень грязным для больших тензоров, и вам нужно запустить операцию tf.assign (), прежде чем значение будет установлено.

В numpy я бы просто сделал что-то вроде

import numpy as np
A = np.array([[1,2,3],[4,5,6],[7,8,9]])
B = np.zeros((7,3))
for i in range(3):
    B[2*i:2*i+3,i] = A[i,:]

Есть ли какой-то похожий или даже более простой способ сделать это в Tensorflow, при котором операция все еще дифференцируема?

Answer 1

Вы можете сделать это с помощью tf.scatter_nd следующим образом (функция работает для TF 1.x и 2.x):

import tensorflow as tf

def make_tensor(a):
    a = tf.convert_to_tensor(a)
    s = tf.shape(a)
    n = s[0]
    m = s[1]
    out_shape = [1 + (m - 1) * n, n]
    r = tf.expand_dims(tf.range(n), 1)
    idx_row = r * (m - 1) + tf.range(m)
    idx_col = tf.tile(r, [1, m])
    return tf.scatter_nd(tf.stack([idx_row, idx_col], axis=-1), a, out_shape)

print(make_tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]]).numpy())
# [[1 0 0]
#  [2 0 0]
#  [3 4 0]
#  [0 5 0]
#  [0 6 7]
#  [0 0 8]
#  [0 0 9]]

РЕДАКТИРОВАТЬ: Если вы хотите, вы можетеНемного измените приведенную выше функцию, чтобы выбрать степень перекрытия между столбцами:

import tensorflow as tf

def make_tensor(a, overlap=1):
    s = tf.shape(a)
    n = s[0]
    m = s[1]
    m2 = m - overlap
    out_shape = [overlap + m2 * n, n]
    r = tf.expand_dims(tf.range(n), 1)
    idx_row = r * m2 + tf.range(m)
    idx_col = tf.tile(r, [1, m])
    return tf.scatter_nd(tf.stack([idx_row, idx_col], axis=-1), a, out_shape)

a = tf.constant([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
print(make_tensor(a).numpy())
# [[1 0 0]
#  [2 0 0]
#  [3 4 0]
#  [0 5 0]
#  [0 6 7]
#  [0 0 8]
#  [0 0 9]]
print(make_tensor(a, overlap=2).numpy())
# [[1 0 0]
#  [2 4 0]
#  [3 5 7]
#  [0 6 8]
#  [0 0 9]]
print(make_tensor(a, overlap=3).numpy())
# [[1 4 7]
#  [2 5 8]
#  [3 6 9]]

Что позволяет случайно не перекрывать или даже отрицательно перекрывать:

print(make_tensor(a, overlap=0).numpy())
# [[1 0 0]
#  [2 0 0]
#  [3 0 0]
#  [0 4 0]
#  [0 5 0]
#  [0 6 0]
#  [0 0 7]
#  [0 0 8]
#  [0 0 9]]
print(make_tensor(a, overlap=-1).numpy())
# [[1 0 0]
#  [2 0 0]
#  [3 0 0]
#  [0 0 0]
#  [0 4 0]
#  [0 5 0]
#  [0 6 0]
#  [0 0 0]
#  [0 0 7]
#  [0 0 8]
#  [0 0 9]]