PROSAGA码农传奇-pytorch-在pytorch Argmax中发生冲突时的索引选择

<div class =“post-text”itemprop =“text”>
  <P>
    
这是一个很好的问题，我自己偶然发现了几次。最简单的答案是，没有任何保证
     <code>
      torch.argmax
    </code>
     （要么
     <code>
      torch.max(x, dim=k)
    </code>
    ，当指定dim时也返回索引）将一致地返回相同的索引。相反，它会回来
    <EM>
      任何有效的索引
    </EM>
     到argmax值，可能是随机的。如
    <a href="https://discuss.pytorch.org/t/function-torch-max-return-indices-inconsistency-between-cup-and-gpu/1890/2" rel="nofollow noreferrer">
      这个帖子在官方论坛上
    </A>
     讨论，这被认为是理想的行为。 （我知道我刚才读过另一个线程，这使得这个更明确，但我再也找不到了）。
  </p>
  <P>
    话虽如此，由于这种行为对我的用例来说是不可接受的，我写了下面的函数，它们会找到最左边和最右边的索引（请注意
     <code>
      condition
    </code>
     是传入的函数对象）：
  </p>
   <pre>
    <code>
      def __consistent_args(input, condition, indices):
    assert len(input.shape) == 2, 'only works for batch x dim tensors along the dim axis'
    mask = condition(input).float() * indices.unsqueeze(0).expand_as(input)
    return torch.argmax(mask, dim=1)

def consistent_find_leftmost(input, condition):
    indices = torch.arange(input.size(1), 0, -1, dtype=torch.float, device=input.device)
    return __consistent_args(input, condition, indices)

def consistent_find_rightmost(input, condition):
    indices = torch.arange(0, input.size(1), 1, dtype=torch.float, device=input.device)
    return __consistent_args(input, condition, indices)

# one example:
consistent_find_leftmost(torch.arange(10).unsqueeze(0), lambda x: x>5)                                                                                                                                     
# will return: 
# tensor([6])

</code>
  </pre>
  <P>
    希望他们能帮忙！ （哦，如果你有更好的实现，请告诉我）
  </p>
</DIV>