softmax(a,axis=0)的用法理解总结

关注

转载最新推荐文章于 2026-05-20 10:11:29 发布 1k 阅读 AI 写同款· GEO 优化 ›

PyTorch中dim参数实战指南：从dim=0到dim=-1的softmax应用详解本文深入解析了PyTorch中softmax函数的dim参数，从张量维度基础讲起，通过三维张量实例详细拆解了dim=0、1、2及dim=-1的计算逻辑与区别。文章结合图像分类、NLP序列模型及注意力机制等实战场景，阐述了如何正确选择dim参数以得到有效的概率分布，并提供了数值稳定性处理及高维张量调试等高级技巧。阅读详情

对于3维度数组总结axis=0 , 1 ,2

axis=0 沿着 axis=0方向 (可以认为是时间的方向)取每个单元对应元素进行计算softmax() //通俗理解就是今天8点钟的对应行对应列的元素a[2][0][0] 与昨天8点钟的对应行对应列的元素a[1][0][0] 及前天8点钟的对应行对应列的元素a[0][0][0] 都取出来，进行计算softmax()

axis=1 沿着 axis=1方向 (行的方向)取每个单元对应元素进行计算softmax()

axis=2 沿着 axis=2方向 (列的方向)取每个单元对应元素进行计算softmax()

import tensorflow as tf
import numpy as np
print("tf.__version__=",tf.__version__)
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)

a1 = np.array([[[1, 2, 3], [2, 4, 6]],[[2, 4, 6], [4, 5, 6]]])
a1 = tf.cast(a1, 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)>
#
s11 = tf.nn.softmax(a1,axis=0)
print(s11)

输出

tf.__version__= 2.0.0
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)
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)
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)
tf.Tensor(
[[[0.26894143 0.11920291 0.04742587]
  [0.11920291 0.26894143 0.5       ]]

 [[0.7310586  0.880797   0.95257413]
  [0.880797   0.7310586  0.5       ]]], shape=(2, 2, 3), dtype=float32)

https://blog.youkuaiyun.com/q2519008/article/details/107086024?utm_medium=distribute.pc_relevant_t0.none-task-blog-BlogCommendFromBaidu-1.control&depth_1-utm_source=distribute.pc_relevant_t0.none-task-blog-BlogCommendFromBaidu-1.control

收录于