PyTorch：如何用小张量的argmax索引大张量并保持顺序？

问题背景 #

我们有两个张量：

• 形状为[2, 1024]的张量b，示例如下：

b = [[2.4, 3.2, 1.5, 6.9, ...],
     [5.9, 3.5, 1.2, 2.0, ...]]

• 形状为[2, 1024, 10]的张量t

需要利用b沿第0维的argmax结果，从t的第0或第1个切片中取出对应行，最终得到形状为[1024, 10]的张量output，且严格保持原有元素顺序。

代码实现 #

PyTorch版本

import torch

# 假设b和t已定义
idx = b.argmax(dim=0)  # 形状 [1024]，每个元素为0或1
output = t[idx, torch.arange(1024), :]  # 形状 [1024, 10]

逻辑说明：

• argmax(dim=0)计算b每一列（沿第0维）最大值所在的行索引，得到长度为1024的索引数组。
• torch.arange(1024)生成列位置索引，配合idx精准定位t中需要提取的元素，确保每个位置i取t[idx[i], i, :]，完全匹配需求。

TensorFlow版本

import tensorflow as tf

# 假设b和t已定义
idx = tf.argmax(b, axis=0)  # 形状 [1024]，每个元素为0或1
# 构造索引对：(行索引, 列索引)
indices = tf.stack([idx, tf.range(1024)], axis=1)
output = tf.gather_nd(t, indices)  # 形状 [1024, 10]

逻辑说明：

• tf.argmax(b, axis=0)得到每列最大值的行索引。
• tf.stack将行索引和列位置索引配对成[1024, 2]的索引矩阵，通过tf.gather_nd从t中提取对应位置的10维向量，最终得到目标形状的张量。

内容的提问来源于stack exchange，提问作者Charlie