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