如何在张量流中的两个选定维度中执行 tf.nn.softmax？

Question

我想为形状为 (batch_size=?, height, width, channel) 的选定二维张量实现 tf.nn.softmax()。

但是tf.nn.softmax()好像不能同时接收2轴。使用 tf.softmax(tensor, axis=[1, 2]) 会增加 tensorflow 中的轴误差。

如何在矢量化模式下优雅地实现它？谢谢：D

Answer 1

你可以做到

array = np.random.rand(1, 2, 2, 1)
s1 = tf.nn.softmax(array, axis=1)
s2 = tf.nn.softmax(array, axis=2)
rs = tf.reduce_sum([s1, s2], 0)

这将 return 个与初始数组形状相同的张量

Answer 2

我不会一次传递两个维度，而是首先相应地重塑输入，例如：

array = tf.constant([[1., 2.], [3., 4.]])

tf.nn.softmax(array, axis=0) # Calculate for axis 0
tf.nn.softmax(array, axis=1) # Calculate for axis 1

tf.nn.softmax(tf.reshape(array, [-1])) # Calculate for both axes

Answer 3

可以用keras激活函数来完成：

# logits has shape [BS, H, W, CH]
prob = tf.keras.activations.softmax(logits, axis=[1, 2])
# prob has shape [BS, H, W, CH]
check = tf.reduce_sum(prob, axis=[1, 2])
# check is tensor of ones with shape [BS, CH]

如何在张量流中的两个选定维度中执行 tf.nn.softmax？

How to perform tf.nn.softmax in two selected dimension in tensorflow?

tensorflow

softmax