Pytorch 二维矩阵中关于dim使用的样例

import torch
a = torch.tensor([[1.0, 2.0, 3.0],
                  [2.0, 2.0, 2.0],
                  [3.0, 2.0, 1.0]])
print(a.shape)
print(a)
"""
torch.Size([3, 3])
tensor([[ 1.,  2.,  3.],
        [ 2.,  2.,  2.],
        [ 3.,  2.,  1.]])
"""

使用softmax函数，当dim=1时，矩阵中每一行的和为1：

import torch.nn.functional as F
attention = F.softmax(a,dim=1)
print(attention)
"""
tensor([[ 0.0900,  0.2447,  0.6652],
        [ 0.3333,  0.3333,  0.3333],
        [ 0.6652,  0.2447,  0.0900]])
"""

dim=0时，矩阵中每一列的和为1：

attention_1 = F.softmax(a,dim=0)
print(attention_1)
"""
tensor([[ 0.0900,  0.3333,  0.6652],
        [ 0.2447,  0.3333,  0.2447],
        [ 0.6652,  0.3333,  0.0900]])
"""

按照我目前的理解，当维度大于2时，dim参数就是对指定的维度进行处理了，二维矩阵的处理总是抱有疑惑，故本篇博客以一个样例来说明。

分享一个使用方括号的角度来对维度进行解读的博客：
【Pytorch】Pytorch 中的 dim