我有一个4D张量,我想得到最后二维的argmaxtorch.argmax
只接受整数作为“dim”参数,而不接受元组
我怎样才能做到这一点
这是我的想法,但我不知道如何匹配我的两个“指数”张量的维数original_array
是形状[1,512,37,59]
max_vals, indices_r = torch.max(original_array, dim=2)
max_vals, indices_c = torch.max(max_vals, dim=2)
indices = torch.hstack((indices_r, indices_c))
正如其他人提到的,最好将最后两个维度展平并应用argmax
。。您将获得最大值的线性索引。如果需要最大值的二维索引,可以使用列数“取消设置”索引
这将为您提供一个
(1, 512, 2)
,其中最后两个dim包含2D坐标您可以使用^{} 展平最后两个维度并在其上应用^{} :
相关问题 更多 >
编程相关推荐