在 tf.nn.top_k 中合并 torch.topk 的 dim

Question

Pytorch 提供 torch.topk(input, k, dim=None, largest=True, sorted=True) 函数来计算给定 input 张量沿给定维度 dim.

的最大元素 k

我有一个形状为 (64, 128, 512) 的张量，我按以下方式使用 torch.topk-

reduce = input.topk(k, dim=1).values

我发现类似的 tensorflow 实现如下 - tf.nn.top_k(input, k=1, sorted=True, name=None)。

我的问题是如何在tf.nn.top_k中加入dim=1参数，从而得到与pytorch计算的相同形状的张量？

Answer 1

我同意@jodag，你必须转置或重塑你的张量，因为tf.math.top_k总是在最后一个维度上工作。

你还可以做的是首先沿着某个维度获取张量中的所有最大值，然后从该张量中获取最高的 k 值：

import tensorflow as tf
tf.random.set_seed(2)

k = 3
tensor = tf.random.uniform((2, 4, 6), maxval=10, dtype=tf.int32)
max_tensor = tf.reduce_max(tensor, axis=1)
k_max_tensor = tf.math.top_k(max_tensor, k=k, sorted=True).values

print('Original tensor --> ', tensor)
print('Max tensor --> ', max_tensor)
print('K-Max tensor --> ', k_max_tensor)
print('Unique K-Max tensor', tf.unique(tf.reshape(k_max_tensor, (tf.math.reduce_prod(tf.shape(k_max_tensor)), ))).y)

Original tensor -->  tf.Tensor(
[[[1 6 2 7 3 6]
  [7 5 1 1 0 6]
  [9 1 3 9 1 4]
  [6 0 6 2 4 0]]

 [[4 6 8 2 4 7]
  [5 0 8 2 8 9]
  [0 2 0 0 9 8]
  [9 3 8 9 0 6]]], shape=(2, 4, 6), dtype=int32)
Max tensor -->  tf.Tensor(
[[9 6 6 9 4 6]
 [9 6 8 9 9 9]], shape=(2, 6), dtype=int32)
K-Max tensor -->  tf.Tensor(
[[9 9 6]
 [9 9 9]], shape=(2, 3), dtype=int32)
Unique K-Max tensor tf.Tensor([9 6], shape=(2,), dtype=int32)

在 tf.nn.top_k 中合并 torch.topk 的 dim

Incorporating dim of torch.topk in tf.nn.top_k

python

tensorflow

pytorch