如何用Keras后端实现2D张量的逐行指定索引（GPU兼容）？

要实现你描述的需求——用形状为(n,)的整数张量作为每行的索引，从(n,k)的2D张量中提取对应位置的元素，并且适配GPU运行，我们可以利用Keras后端的张量操作函数实现完全向量化计算（避免列表推导式这种CPU串行的方式）。

核心思路 #

我们需要构造行-列索引对，然后用Keras后端的gather_nd函数一次性提取所有目标元素，具体步骤：

• 生成行索引序列：从0到n-1的整数张量，对应2D张量的每一行
• 将行索引与输入的列索引张量拼接成形状为(n,2)的索引矩阵，每一行代表一个(行号, 列号)的坐标
• 用gather_nd根据这个索引矩阵从2D张量中取值

实现代码 #

import keras.backend as K

def gather_rowwise(matrix, col_indices):
    """
    从形状为(n,k)的matrix张量中，按col_indices（形状(n,)）的每个元素作为列索引，提取每行对应位置的元素
    返回形状为(n,)的张量
    """
    # 获取矩阵的行数n
    n = K.shape(matrix)[0]
    # 生成行索引：[0, 1, 2, ..., n-1]
    row_indices = K.arange(n, dtype=K.dtype(col_indices))
    # 拼接行索引和列索引，得到(n,2)的索引对矩阵
    index_pairs = K.stack([row_indices, col_indices], axis=1)
    # 根据索引对提取元素
    return K.gather_nd(matrix, index_pairs)

测试验证 #

用你给出的示例数据验证函数效果：

import numpy as np
from keras.layers import Input
from keras.models import Model

# 测试数据
col_indices = np.array([1,2,0,0])
matrix = np.array([[1,2,3],[2,3,4],[2,3,1],[3,2,1]])

# 创建输入张量（匹配数据形状）
matrix_input = Input(shape=(3,), batch_shape=(4,))
indices_input = Input(shape=(), batch_shape=(4,), dtype='int32')

# 应用自定义函数
output_tensor = gather_rowwise(matrix_input, indices_input)

# 构建模型并预测
model = Model(inputs=[matrix_input, indices_input], outputs=output_tensor)
result = model.predict([matrix, col_indices])

print(result)  # 输出: [2. 4. 2. 3.]，和列表推导式的结果完全一致

为什么适配GPU？ #

这个实现全程使用Keras后端的张量操作，所有计算都会被转换为TensorFlow（或其他后端）的计算图，能够充分利用GPU的并行计算能力，不会像列表推导式那样在CPU上逐元素循环，完美适配GPU运行场景。

如果你的场景是批量输入（比如形状为(batch_size, n, k)的矩阵和(batch_size, n)的索引），只需要稍作调整索引的构造逻辑，比如增加batch维度的处理，同样可以实现向量化计算。

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