擅长:python、mysql、java
<p>正如其他人提到的,最好将最后两个维度展平并应用argmax</p>
<pre><code>original_array = torch.rand(1, 512, 37, 59)
original_flatten = original_array.view(1, 512, -1)
_, max_ind = original_flatten.max(-1)
</code></pre>
<p>。。您将获得最大值的线性索引。如果需要最大值的二维索引,可以使用列数“取消设置”索引</p>
<pre><code># 59 is the number of columns for the (37, 59) part
torch.stack([max_ind // 59, max_ind % 59], -1)
</code></pre>
<p>这将为您提供一个<code>(1, 512, 2)</code>,其中最后两个dim包含2D坐标</p>