Pytorch和Python的坑

张量

import torch
a = torch.tensor(  //一个三维张量
    [
        [
            [1,5,62,54],
            [2,6,2,6],
            [2,65,2,6]
         ]
         ,
         [
            [1,5,62,54],
            [2,6,2,6],
            [2,65,2,6]]
         ]
 )
c , d = torch.max(a, 2)#max就是按照特定的维度取max。返回两个张量，一个是最大值，一个是最大值的下标
print(c)
print(d) 

#维度为2的结果
tensor([[[ 1,  5, 62, 54],
         [ 2,  6,  2,  6],
         [ 2, 65,  2,  6]],

        [[ 1,  5, 62, 54],
         [ 2,  6,  2,  6],
         [ 2, 65,  2,  6]]])
tensor([[62,  6, 65],
        [62,  6, 65]])
tensor([[2, 3, 1],
        [2, 3, 1]])

#维度为1的结果
tensor([[[ 1,  5, 62, 54],
         [ 2,  6,  2,  6],
         [ 2, 65,  2,  6]],

        [[ 1,  5, 62, 54],
         [ 2,  6,  2,  6],
         [ 2, 65,  2,  6]]])
tensor([[ 2, 65, 62, 54],
        [ 2, 65, 62, 54]])
tensor([[2, 2, 0, 0],
        [2, 2, 0, 0]])

#维度为0的结果
tensor([[[ 1,  5, 62, 54],
         [ 2,  6,  2,  6],
         [ 2, 65,  2,  6]],

        [[ 1,  5, 62, 54],
         [ 2,  6,  2,  6],
         [ 2, 65,  2,  6]]])
tensor([[ 1,  5, 62, 54],
        [ 2,  6,  2,  6],
        [ 2, 65,  2,  6]])
tensor([[1, 1, 1, 1],
        [1, 1, 1, 1],
        [1, 1, 1, 1]])

torch.sum

a = torch.tensor(  # 一个三维张量
        [
            [1,5,62,54],
            [2,6,2,6],
            [2,65,2,6]
         ]
 )

sum_fg = torch.sum((a == 2).int(), 1)
print(sum_fg)

下面是结果

tensor([0, 2, 2])