softmax中axis参数

最新推荐文章于 2025-07-22 00:16:12 发布

原创最新推荐文章于 2025-07-22 00:16:12 发布 · 2.7k 阅读

9 ·

CC 4.0 BY-SA版权

文章标签：

#softmax

DL 专栏收录该内容

12 篇文章

订阅专栏

本文通过实例详细解析了在TensorFlow中使用Softmax函数时，轴参数(axis)的不同设置如何影响计算结果，对比了在二维和三维数组上的具体差异。

从caffe中我们看到softmax有下面这些参数

// Message that stores parameters used by SoftmaxLayer, SoftmaxWithLossLayer
message SoftmaxParameter {
  enum Engine {
    DEFAULT = 0;
    CAFFE = 1;
    CUDNN = 2;
  }
  optional Engine engine = 1 [default = DEFAULT];

  // The axis along which to perform the softmax -- may be negative to index
  // from the end (e.g., -1 for the last axis).
  // Any other axes will be evaluated as independent softmaxes.
  optional int32 axis = 2 [default = 1];
}

一般来说axis也不需要修改，默认设为1，即在c上做计算。
那么设置不同的axis，结果有什么不同？我们举个例子一目了然；

import tensorflow as tf
import numpy as np

a = np.array([[1, 2, 3], [1, 2, 3]])
a = tf.cast(a, tf.float32)
#>>> a
#tf.Tensor: shape=(2, 3), dtype=float32, numpy=
#array([[1., 2., 3.],
#       [1., 2., 3.]], dtype=float32)>
# 

s1 = tf.nn.softmax(a,axis=0)
print(s1)
#tf.Tensor(
#[[0.5 0.5 0.5]
#[0.5 0.5 0.5]], shape=(2, 3), dtype=float32)

s2 = tf.nn.softmax(a,axis=1)
print(s2)
#tf.Tensor(
#[[0.09003057 0.24472848 0.66524094]
#[0.09003057 0.24472848 0.66524094]], shape=(2, 3), dtype=float32)

我们来看看计算过程：
axis = 0时（表示纵轴，方向从上到下）
在这里插入图片描述

axis = 1时（表示横轴，方向从左到右）
在这里插入图片描述

再举一个三维数组深入理解一下

import tensorflow as tf
import numpy as np

a = np.array([[[1, 2, 3], [1, 2, 3]],[[4, 5, 6], [4, 5, 6]]])
a = tf.cast(a, tf.float32)
#>>> a
#<tf.Tensor: shape=(2, 2, 3), dtype=float32, numpy=
#array([[[1., 2., 3.],
#        [1., 2., 3.]],
#
#       [[4., 5., 6.],
#        [4., 5., 6.]]], dtype=float32)>
#
s1 = tf.nn.softmax(a,axis=0)
print(s1)
#tf.Tensor(
#[[[0.04742587 0.04742587 0.04742587]
#  [0.04742587 0.04742587 0.04742587]]
#
# [[0.95257413 0.95257413 0.95257413]
#  [0.95257413 0.95257413 0.95257413]]], shape=(2, 2, 3), dtype=float32)

s2 = tf.nn.softmax(a,axis=1)
print(s2)
#tf.Tensor(
#[[[0.5 0.5 0.5]
#  [0.5 0.5 0.5]]
#
# [[0.5 0.5 0.5]
#  [0.5 0.5 0.5]]], shape=(2, 2, 3), dtype=float32)

s3 = tf.nn.softmax(a,axis=2)
print(s3)
#tf.Tensor(
#[[[0.09003057 0.24472848 0.66524094]
#  [0.09003057 0.24472848 0.66524094]]
#
# [[0.09003057 0.24472848 0.66524094]
#  [0.09003057 0.24472848 0.66524094]]], shape=(2, 2, 3), dtype=float32)