TensorFlow中从Variable提取指定元素组成新矩阵的方法

陆景行

发布时间：2026-10-13 | 6人浏览过

原创

本文介绍如何在TensorFlow中从Variable类型的矩阵中提取指定元素并组成新的Tensor。通过一维数组和二维矩阵示例，详细演示了如何使用tf.gather_nd函数配合索引矩阵提取单个或批量元素，并推广至n维矩阵的处理方式。

TensorFlow Variable提取元素方法概述

在神经网络计算过程中，经常需要提取矩阵中的特定元素进行独立计算（例如MLE、Attention等机制）。本文将演示如何在TensorFlow的Variable类型中实现这一操作。

一维数组的元素提取

首先假设Variable为一个一维数组A：


import numpy as np

import tensorflow as tf

a = np.array([1, 2, 3, 4, 5, 6, 7, 8])

A = tf.Variable(a)

我们将需要提取的元素索引存储到变量B中。若仅需提取数组A中的单个元素，B的定义如下：


b = np.array([3])

B = tf.placeholder(dtype=tf.int32, shape=[1])

由于索引坐标仅有一维，因此shape设为1。

提取元素并将其组合为Tensor C的操作如下：


C = tf.gather_nd(A, B)

运行上述代码：


init = tf.global_variables_initializer()

with tf.Session() as sess:
  init.run()
  feed_dict = {B: b}
  result = sess.run([C], feed_dict=feed_dict)
  print result

输出结果为：

[4]

若需从一维数组中提取多个元素，则需将每个目标元素的索引单独列为一行：


b = np.array([[3], [2], [5], [0]])

B = tf.placeholder(dtype=tf.int32, shape=[4, 1])

此时，由于要从一维数组中索引4个数，shape应设为[4, 1]。

再次运行代码得到：


[4 3 6 1]

////////////////////////////////////////////////////////////////////////////////////华丽丽的分割线

二维矩阵的元素提取

假设Variable为一个二维矩阵A：


a = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])

A = tf.Variable(a)

首先提取A中的一个元素，需给定该元素的行列坐标，并存入B中：


b = np.array([2，3])

B = tf.placeholder(dtype=tf.int32, shape=[2])

注意，由于输入的索引坐标变为二维，shape也随之变为2。

提取元素并组合为Tensor C：


C = tf.gather_nd(A, B)

运行代码：


init = tf.global_variables_initializer()

with tf.Session() as sess:
  init.run()
  feed_dict = {B: b}
  result = sess.run([C], feed_dict=feed_dict)
  print result

输出结果为：


[12]

同理，若要从二维矩阵中提取多个元素，需将每个目标元素的索引单独列为一行：


b = np.array([[2, 3], [1, 0], [2, 2], [0, 1]])

B = tf.placeholder(dtype=tf.int32, shape=[4, 2])

此时，由于要从二维矩阵中索引4个数，shape应设为[4, 2]。

再次运行代码得到：


[12 5 11 2]

////////////////////////////////////////////////////////////////////////////////////华丽丽的分割线

n维矩阵的通用处理方法

推广至n维矩阵的情况：

假设A是Variable类型的n维矩阵，若要提取其中的m个元素，首先需将每个元素的索引坐标表示为列表形式：


index = [x1, x2, x3, ..., xn]

其中xj代表该元素在n维矩阵中第j维的位置。

其次，每个坐标需单独占据索引矩阵的一行：


index_matrix = [[x11, x12, x13, ..., x1n],

               [x21, x22, x23, ..., x2n],

               [x31, x32, x33, ..., x3n],

               .......................................,

               [xm1, xm2, xm3, ..., xmn]]

最后使用tf.gather_nd()函数进行替换即可：