python 从tensorflow 中返回k个最小元素

hmmo2u0o  于 2022-12-17  发布在  Python
关注(0)|答案(1)|浏览(144)

我有一个Tensor,在tensorflow 中的形状是(16, 512, 4096),我想从Tensor中计算出k的最小元素。
请注意,我可以使用以下代码片段在pytorch中获取它-

#inputs.shape (16L, 512L, 4096L)
dists, inputs_idx = torch.topk(inputs, 64, 2, largest=False, sorted=False)
#dists.shape (16L, 512L, 64L), inputs_idx.shape (16L, 512L, 64L)

有什么解决办法吗?

epggiuax

epggiuax1#

因为torch.topk可以用来得到k最大的元素,所以你可以对这些值取反,执行tne操作,然后再次对它们取反以得到值:

-torch.topk(-x, ...)

相关问题