[PyTorch] 如何判定运算维度

实际上无论是几维，方法都是一样。假设以 torch.softmax() 为例：

# 下面运行结果所使用的代码
import torch
import numpy as np
z = np.arange(1, 33).reshape((2, 2, 2, 4))
z = torch.tensor(z, dtype=torch.float32)
# 为了使各元素softmax的结果相差不至于过大，这里简单处理一下，但是与原张量的计算结果并不等价
z = z ** 0.2
torch.softmax(z, dim=0) 

一. 三维

为了方便查阅，直接放个图，方法实际上和四维是一样的。

二. 四维

假设有四维tensor：(B, C, H, W)，具体为(2, 2, 2, 4)

$$\begin{align} [[[&[\space1, \quad 2, \quad 3, \quad 4], \\ &[\space5, \quad 6, \quad 7, \quad 8]], \\ [&[\space9, \space\space\space 10, \space\space 11, \space\space 12], \\ &[13, \space\space 14, \space\space 15, \space\space 16]]], \\ [&[17, \space\space 18, \space\space 19, \space\space 20], \\ &[21, \space\space 22, \space\space 23, \space\space 24]], \\ [&[25, \space\space 26, \space\space 27, \space\space 28], \\ &[29, \space\space 30, \space\space 31, \space\space 32]]]] \end{align}$$

---

1. 当dim=0或dim=-4时

最外层 $[\space]$ 的元素进行运算，也就是最外层 $[\space]$ 中用逗号隔开的所有对应元素相运算。例如向量中的1、17进行运算，9和25运算。

---

2. 当dim=1或dim=-3时

同一 batch 的最外层 $[\space]$ 的所有对应元素进行运算。例如向量中的1和9进行运算，17和25运算。

---

3. 当dim=2或dim=-2时

同一 Height 的最外层 $[\space]$ 的所有对应元素进行运算。例如向量中的1和5进行运算，9和13进行运算，17和21进行运算，25和29进行运算。

---

4. 当dim=3或dim=-1时

同一 W 的最外层 $[\space]$ 的所有对应元素进行运算。例如向量中的1、2、3、4进行运算，5、6、7、8进行运算，……，29、30、31、32进行运算。

三. 高维

类比四维的情况……

【注】可能我表达确实不行，如看不懂，可看看这篇文章：https://zhuanlan.zhihu.com/p/525276061