pytorch中的eq()/gt()/lt()/t()/topk()函数

a = torch.Tensor([0.1,0.8])
b = torch.Tensor([0,1])
c = torch.Tensor([[2,1],[3,4]])

print('a等于b:',a.eq(b))
print('a大于b:',a.gt(b))
print('a小于b:',a.lt(b))
print('c:',c)
print('c转置:',c.t())
#print结果为
a等于b: tensor([False, False])
a大于b: tensor([ True, False])
a小于b: tensor([False,  True])
c: tensor([[2., 1.],
        [3., 4.]])
c转置: tensor([[2., 3.],
        [1., 4.]])

topk()函数

>>> a=torch.randn((4,6))
>>> print(a)
tensor([[ 0.7042,  0.2533,  1.1596, -0.7436,  0.5264,  0.2085],
        [ 0.2641,  0.9683,  0.4469, -1.9215, -0.7564,  1.1776],
        [ 1.0520, -1.6003, -0.8634,  1.7596, -0.8464,  0.7166],
        [-0.0492, -0.7746,  1.2592, -0.8273,  0.1266,  1.0450]])
>>> maxk=max((1,3))
>>> _, pred=a.topk(maxk,1,True,True)
>>> print(_)
tensor([[1.1596, 0.7042, 0.5264],
        [1.1776, 0.9683, 0.4469],
        [1.7596, 1.0520, 0.7166],
        [1.2592, 1.0450, 0.1266]])
>>> print(pred)
tensor([[2, 0, 4],
        [5, 1, 2],
        [3, 0, 5],
        [2, 5, 4]])
>>> _, pred=a.topk(1,1,True,True)
>>> print(_)
tensor([[1.1596],
        [1.1776],
        [1.7596],
        [1.2592]])
>>> print(pred)
tensor([[2],
        [5],
        [3],
        [2]])

>>> _, pred=a.topk(1,0,True,True)
>>> print(_)
tensor([[1.0520, 0.9683, 1.2592, 1.7596, 0.5264, 1.1776]])
>>> print(pred)
tensor([[2, 1, 3, 2, 0, 1]])

_返回的是前maxk个最大值，pred返回的是对应的index，第二个参数是指定维度dim=0，按行取，dim=1，按列取。
max()函数也可以有相同的功能，但需要设置keepdim=True，避免降维。

_, indices_max = a.max(dim=1, keepdim=True)