torch中函数的一些用法总结(不断更新)

torch.gather

b = torch.Tensor([[1,2,3],[4,5,6]])
print(b)
index_1 = torch.LongTensor([[0,1],[2,0]])
index_2 = torch.LongTensor([[0,1,1],[0,0,0]])
print(torch.gather(b, dim=1, index=index_1))
print(torch.gather(b, dim=0, index=index_2))
b:
tensor([[1., 2., 3.],
        [4., 5., 6.]])
index_1:
tensor([[0, 1],
        [2, 0]]) 
        
index_2:
tensor([[0, 1, 1],
        [0, 0, 0]])
print(torch.gather(b, dim=1, index=index_1))
tensor([[1., 2.],
        [6., 4.]])

print(torch.gather(b, dim=0, index=index_2))     
tensor([[1., 5., 6.],
        [1., 2., 3.]])
  • 如果dim = 1:index_1中元素为对应b中行元素的索引。
  • 如果dim = 2:index_2第一列元素(0,0)对应b中的第一列,所以取出(1,1)元素,index_2第二列元素(1,0)对应b中的第二列,所以取出(5,2)元素,index_2第三列元素(1,0)对应b中的第一列,所以取出(6,3)元素.

torch.eq(predict_ labels, labels)

式中predict_ labels与labels是两个大小相同的tensor,而torch.eq()函数就是用来比较对应位置数字,相同则为True,否则为False,输出与那两个tensor大小相同,并且其中只有True和False。

predict_labels = torch.LongTensor([0,1, 2, 3 ,4])
labels =  torch.LongTensor([4,3,2,1,4])

torch.eq(predict_labels, labels)
Out[15]: tensor([False, False,  True, False,  True])

torch.eq(predict_labels, labels).sum()
Out[16]: tensor(2)

torch.eq(predict_labels, labels).sum().item()
Out[17]: 2

猜你喜欢

转载自blog.csdn.net/weixin_54546190/article/details/126585016