Tensorflow.js tf.gatherND() 函数介绍

Question

Tensorflow.js 中的 tf.gatherND() ​ 函数通过提供索引来从张量中收集元素。它支持索引多个元素，包括多个维度和不同形状的张量。tf.gatherND() 函数也可以用于子集采样，其中张量中的一部分元素被采样。

语法

以下是 tf.gatherND() 函数的语法：

tf.gatherND(x, indices)

• x: 张量，要从中抽取元素的张量。
• indices: 张量，要在 x 中选择元素的索引。indices 张量的最后一个维度必须与 x 的最后一个维度匹配，并且它们必须具有不同的维度大小。最后一个维度可以是任何大小。

示例

const x = tf.tensor3d([[[0, 1], [2, 3]], [[4, 5], [6, 7]]]);
const indices = tf.tensor2d([[1, 0], [0, 1]]);
const result = tf.gatherND(x, indices);

result.print();
// output: 
// [[2, 1],
//  [4, 7]]

在上面的示例中，我们使用 tf.tensor3d() 创建了一个形状为[2, 2, 2]的 3D 张量 x。然后，我们使用 tf.tensor2d() 创建一个形状为[2, 2]的 2D 张量 indices，该张量表示要从 x 中选择的元素的索引。我们使用 tf.gatherND() 函数来选择这些元素并存储在结果张量中，并使用 print() 函数在控制台中打印结果。

注意事项

• indices 的形状必须是 [s1,...,sn, r, q]，其中 s1,...,sn 是在 x 上的形状，r 是最后一个维度大小，q 可以是任何大小。
• tf.gatherND() 函数返回的张量的形状与 indices 的形状相同，其中最后一个维度被替换为 x 张量中的相应维度大小。

总结

tf.gatherND() 函数是一个有用的函数，可以根据索引从张量中选择元素。这个函数支持多个维度和不同形状的张量，可以用于子集采样。本文介绍了 tf.gatherND() 函数的用法，示例和注意事项，希望可以帮助您更好地使用 Tensorflow.js。