Tensorflow：跨批次的 2x2 矩阵中的最大值索引

Question

如果我有一批矩阵，我的矩阵是形状 (?, 600, 600)，我将如何检索批处理中每个矩阵中最大值的行和列索引？这样我的行和列返回矩阵都是形状（？）（行返回矩阵具有批处理中每个示例的最大值行的索引，并且对于 col 返回矩阵类似）。

谢谢！

Answer 1

你可以重塑 + argmax。比如：

x = tf.reshape(matrix, [tf.shape(matrix, 0), -1])
indices = tf.argmax(x, axis=1)  # this gives you indices from 0 to 600^2
col_indices = indices / 600
row_indices = indices % 600
final_indices = tf.transpose(tf.stack(col_indices, row_indices))