torch.gather与torch.scatter

2022/1/23 6:07:52

本文主要是介绍torch.gather与torch.scatter,对大家解决编程问题具有一定的参考价值,需要的程序猿们随着小编来一起学习吧!

前言

我觉得pytorch最难的两个api之一就是gather和scatter了,但是在项目中又经常出现,OMG…
还是想办法弄懂吧



torch.gather

input = torch.randn((2, 3, 4))
tensor([[[ 0.1815, -0.7603,  0.3465,  0.2593],
         [ 0.5841,  1.2581,  0.1583, -0.7283],
         [-0.0093,  0.6214, -0.3391, -2.4973]],

        [[-0.8153, -0.3802, -1.9930, -2.1965],
         [-0.8470,  0.5285,  0.2684, -0.2982],
         [-0.6088, -0.6239,  1.1218, -0.1660]]])

index = torch.tensor([[[0, 0, 1]]]).long()
torch.gather(input, 0, index)
tensor([[[ 0.1815, -0.7603, -1.9930]]])

index = torch.tensor([[[0, 2, 1]]]).long()
torch.gather(input, 1, index)
tensor([[[0.1815, 0.6214, 0.1583]]])

index = torch.tensor([[[0, 2, 1]]]).long()
torch.gather(input, 2, index)
tensor([[[ 0.1815,  0.3465, -0.7603]]])





torch.scatter



这篇关于torch.gather与torch.scatter的文章就介绍到这儿,希望我们推荐的文章对大家有所帮助,也希望大家多多支持为之网!


扫一扫关注最新编程教程