向量化两个3D张量的按元素乘积

CKM 发表于 Dev

知识管理

有没有一种方法可以对以下代码进行矢量化处理，以便完全删除循环？

x = tf.constant([[[1,2,3],[2,3,4]],[[1,2,3],[2,3,5]]])
t=tf.eye(x.shape[1])[:,:,None]
for i in range(x.shape[0]):
    out = tf.multiply(t,x[i].numpy())
    out=tf.reshape(out, shape=(out.shape[0], out.shape[-1]*out.shape[-2]))
    print(out)

简而言之：如何将3D张量乘以3D张量的每个元素？就我而言：3个D张量是：

tf.Tensor(
[[[1.]
  [0.]]

 [[0.]
  [1.]]], shape=(2, 2, 1), dtype=float32)

和

tf.Tensor(
[[[1 2 3]
  [2 3 4]]

 [[1 2 3]
  [2 3 5]]], shape=(2, 2, 3), dtype=int32)

预期输出：以下2个张量合并在一起，形状为2 * 2 * 6。

tf.Tensor(
[[1. 2. 3. 0. 0. 0.]
 [0. 0. 0. 2. 3. 4.]], shape=(2, 6), dtype=float32)
tf.Tensor(
[[1. 2. 3. 0. 0. 0.]
 [0. 0. 0. 2. 3. 5.]], shape=(2, 6), dtype=float32)

Jdehesa

这是您如何获得该结果的方法：

import tensorflow as tf

x = tf.constant([[[1, 2, 3], [2, 3, 4]],
                 [[1, 2, 3], [2, 3, 5]]], dtype=tf.float32)
t = tf.eye(tf.shape(x)[1], dtype=x.dtype)
# Add one dimension to x and one dimension to t
xt = tf.expand_dims(x, 1) * tf.expand_dims(t, 2)
# Reshape
result = tf.reshape(xt, (tf.shape(x)[0], tf.shape(x)[1], -1))
print(result.numpy())
# [[[1. 2. 3. 0. 0. 0.]
#   [0. 0. 0. 2. 3. 4.]]
#
#  [[1. 2. 3. 0. 0. 0.]
#   [0. 0. 0. 2. 3. 5.]]]

本文收集自互联网，转载请注明来源。

如有侵权，请联系 [email protected] 删除。