为什么tensorflow对softmax函数使用'dim'参数？

Question

为什么tensorflow对softmax函数使用'dim'参数？我们可以使用什么样的张量作为输入？

Answer 1

tf.nn.softmax在输入中接受通用的非空张量。

您可以决定在您想要的每个维度上应用softmax。

通常，softmax应用于输入张量的最后一个维度（这是默认行为）。这是因为通常softmax应用于神经网络输出，通常是一个形状为[batch_size, num_classes]的张量。

但是，您可以决定将softmax应用于形状为[batch_size, num_classes, 2, 1]的张量，并仅在张量的第二维上计算softmax：tf.nn.softmax(tensor, axis=1)