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

时间:2019-03-16 04:10:28

标签: tensorflow softmax

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

但是tf.nn.softmax()似乎不可能同时接收2个轴。使用tf.softmax(tensor, axis=[1, 2])将增加张量流中的轴误差。

如何才能在矢量化模式下实现优雅? :x

3 个答案:

答案 0 :(得分:0)

您可以

from_request

这将返回与初始数组相同形状的张量

答案 1 :(得分:0)

我不会一次通过两个尺寸,而是先相应地调整输入的形状,例如:

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

答案 2 :(得分:0)

可以使用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]