Question

我有一个数组。我想根据此数组的最后一个维度中的值创建一个掩码。在Numpy，我能做到：

import numpy as np
room = np.array([
    [[0, 0, 1], [1, 0, 0], [1, 0, 0]],
    [[1, 0, 0], [0, 0, 1], [1, 0, 0]],
    [[1, 0, 0], [1, 0, 0], [0, 0, 1]]
])
mask = np.apply_along_axis(lambda x: [1, 1, 1] if (x == [0, 0, 1]).all() else [0, 0, 0], axis=-1, arr=room)
result = mask * room
print(result)

在上面的代码中，room是一个基于我创建掩码的（3,3,3）数组。创建的掩码也是一个（3,3,3）数组，它将用于与其他数组相乘以掩盖不需要的元素。

但是我在使用Tensorflow实现同样的问题时遇到了问题。我试过以下代码，

room = tf.constant([
    [[0, 0, 1], [1, 0, 0], [1, 0, 0]],
    [[1, 0, 0], [0, 0, 1], [1, 0, 0]],
    [[1, 0, 0], [1, 0, 0], [0, 0, 1]]
])
room = tf.reshape(room, shape=(9, -1))
mask = tf.map_fn(lambda x: [1, 1, 1] if x == [0, 0, 1] else [0, 0, 0], room)

但它以下列错误结束：

ValueError: The two structures don't have the same number of elements. First structure: <dtype: 'int32'>, second structure: [0, 0, 0].

Answer 1

map_fn的参数dtype允许指定输出的形状，如果它与x的形状不同。

然而，这不是问题。

您将python条件混合到tensorflow操作中：python操作（如if）在图形外部执行，而您想要定义执行所需操作的图形。

让我们深入研究你的问题：

您想要沿第一维展开room变量：map_fn就可以了。
您想检查当前行（3个元素的向量）是否等于[0, 0, 1]。

为此，您需要使用张量流条件tf.cond(pred, true_fn, false_fn)。

请注意pred必须是标量。因此，如果当前row等于您想要的行并将结果减少为单个标量，则仅使用张量流操作进行检查。

这是真的，只需返回常量值[1,1,1]，否则[0,0,0]。

mask = tf.map_fn(lambda row: tf.cond(
    tf.equal(
        tf.reduce_prod(tf.cast(tf.equal(row, tf.constant([0,0,1])), tf.int32)), 1),
    lambda: tf.constant([1,1,1]),
    lambda: tf.constant([0,0,0])), room)

Tensorflow：如何将函数应用于张量的最后一个维度

1 个答案: