torch.argmax的一些补充

最新推荐文章于 2025-04-03 14:25:48 发布

撒旦即可

最新推荐文章于 2025-04-03 14:25:48 发布

阅读量2.1k

点赞数 3

分类专栏： Pytorch

本文链接：https://blog.youkuaiyun.com/qq_39861441/article/details/108575614

版权

Pytorch 专栏收录该内容

36 篇文章

订阅专栏

torch.argmax是不会向后传梯度，但是被选中的部分还是可以传梯度的

import torch

s=torch.rand(1,3,6,6,requires_grad=True)
d=torch.rand(1,3,6,6,requires_grad=True)
p=torch.argmax(s,dim=1).unsqueeze(1)
q=torch.gather(d,dim=1,index=p)
q=q.sum()
loss=(q-1)*(q-1)
loss.backward()
print(s.grad)
print(d.grad)

output:

None
tensor([[[[ 0.0000,  0.0000,  0.0000, 38.2169,  0.0000,  0.0000],
          [38.2169,  0.0000,  0.0000, 38.2169, 38.2169,  0.0000],
          [ 0.0000, 38.2169,  0.0000,  0.0000,  0.0000,  0.0000],
          [ 0.0000, 38.2169,  0.0000,  0.0000,  0.0000,  0.0000],
          [ 0.0000, 38.2169,  0.0000, 38.2169,  0.0000,  0.0000],
          [38.2169,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000]],

         [[38.2169,  0.0000, 38.2169,  0.0000, 38.2169,  0.0000],
          [ 0.0000, 38.2169,  0.0000,  0.0000,  0.0000,  0.0000],
          [ 0.0000,  0.0000, 38.2169, 38.2169,  0.0000, 38.2169],
          [38.2169,  0.0000,  0.0000, 38.2169, 38.2169, 38.2169],
          [ 0.0000,  0.0000, 38.2169,  0.0000,  0.0000, 38.2169],
          [ 0.0000,  0.0000, 38.2169,  0.0000,  0.0000,  0.0000]],

         [[ 0.0000, 38.2169,  0.0000,  0.0000,  0.0000, 38.2169],
          [ 0.0000,  0.0000, 38.2169,  0.0000,  0.0000, 38.2169],
          [38.2169,  0.0000,  0.0000,  0.0000, 38.2169,  0.0000],
          [ 0.0000,  0.0000, 38.2169,  0.0000,  0.0000,  0.0000],
          [38.2169,  0.0000,  0.0000,  0.0000, 38.2169,  0.0000],
          [ 0.0000, 38.2169,  0.0000, 38.2169, 38.2169, 38.2169]]]])