Операция трансляции по строкам 3D-тензора (Tensorflow) - PullRequest
0 голосов
/ 27 апреля 2018

У меня есть тензор p размерности [BATCH_SIZE, 128] и тензоры A, B обоих измерений [528, 128]. Я хотел бы построить новый тензор p 'размера [BATCH_SIZE, 528], где столбец j определен как:

tf.reduce_prod(self.A[j,:] * p + self.B[j,:], axis=1)

В настоящее время у меня есть реализация этого перебора с использованием цикла for, но это очень медленно. Есть ли способ, которым я могу использовать вещание или что-то еще, чтобы ускорить это?

ps = []
for j in xrange(self.A.shape[0]):
    a = self.A[j,:]
    b = self.B[j,:]
    ps.append(tf.reduce_prod(a * p + b, axis=1))
p = tf.stack(ps, axis=1)

1 Ответ

0 голосов
/ 27 апреля 2018

Попробуйте это:

result = tf.reduce_prod(self.A * p[:, tf.newaxis] + self.B, axis=2)
Добро пожаловать на сайт PullRequest, где вы можете задавать вопросы и получать ответы от других членов сообщества.
...