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

最新推荐文章于 2023-12-03 16:29:47 发布

原创最新推荐文章于 2023-12-03 16:29:47 发布 · 1.7k 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习

deep learning & pytorch 专栏收录该内容

4 篇文章

订阅专栏

本文通过实例详细解析了PyTorch中Softmax函数的使用方法，特别是在不同维度上的应用。展示了如何通过设置dim参数使矩阵行或列的元素和等于1，对理解PyTorch中的张量操作及概率分布转换有重要意义。

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