由其他操作员替换 torch.gather？

Question

我有一个脚本代码，其中x1 和x2 大小为1x68x8x8

 tmp_batch, tmp_channel, tmp_height, tmp_width = x1.size()
 x1 = x1.view(tmp_batch*tmp_channel, -1)        
 max_ids = torch.argmax(x1, 1)            
 max_ids = max_ids.view(-1, 1)
            
 x2 = x2.view(tmp_batch*tmp_channel, -1)
 outputs_x_select = torch.gather(x2, 1, max_ids) # size of 68 x 1

至于上面的代码，当我使用旧的onnx 时，我遇到了torch.gather 的问题。因此，我想找到一个替代解决方案，将toch.gather 替换为其他运算符，但输出与上述代码相同。你能给我一些建议吗？

Answer 1

一种解决方法是使用等效的 numpy 方法。如果您在某处包含 import numpy as np 语句，则可以执行以下操作。

outputs_x_select = torch.Tensor(np.take_along_axis(x2,max_ids,1))

如果这给您带来了与毕业相关的错误，请尝试

outputs_x_select = torch.Tensor(np.take_along_axis(x2.detach(),max_ids,1))

---

没有 numpy 的方法：在这种情况下，max_ids 似乎每行只包含一个条目。因此，我相信以下方法会起作用：

max_ids = torch.argmax(x1, 1) # do not reshape
            
x2 = x2.view(tmp_batch*tmp_channel, -1)
outputs_x_select = x2[torch.arange(tmp_batch*tmp_channel),max_ids]