TensorFlow.js中无需调用Tensor.data()提取布尔张量True值索引的方法

当然有更高效的方案！完全不需要通过Tensor.data()把整个大张量转成JS数组——TensorFlow.js内置的tf.where() API就是专门解决这类问题的，它直接在张量运算层面处理，性能比纯JS遍历数组好太多，尤其适合大张量场景。

为什么你的现有方法性能差？ #

你当前的实现需要把整个张量的所有数据同步到JS数组里，这个过程在处理大张量时会产生巨大的CPU开销（毕竟要把底层张量数据拷贝到JS堆），再加上手动遍历每个元素，自然会耗时超过600ms。

用tf.where()实现高效索引提取 #

tf.where()接收一个布尔张量后，会返回一个2D张量，每一行对应一个true值的坐标索引。如果是1D张量，返回的结果是形如[[0], [2]]的结构，我们可以用tf.squeeze()把它压缩成1D的索引张量，最后只需要把这个小得多的索引张量转成JS数组即可。

示例代码：

// 假设你有一个布尔类型的张量（这里用1D张量举例，高维张量也适用）
const boolTensor = tf.tensor1d([true, false, true, false, true]);

// 提取所有true值的索引，用tf.tidy自动清理临时张量
tf.tidy(() => {
  const indicesTensor = tf.where(boolTensor);
  // 把2D的索引张量压缩成1D（仅适用于1D输入张量）
  const squeezedIndices = indicesTensor.squeeze();
  
  // 仅将索引张量转成JS数组（数据量远小于原始大张量）
  squeezedIndices.data().then(indicesArray => {
    indicesArray.forEach(idx => {
      // 在这里处理每个true值的索引，比如执行你的业务逻辑
      console.log(`找到true值的索引：${idx}`);
    });
  });
});

高维张量的处理方式 #

如果你的布尔张量是高维的（比如2D图像张量），tf.where()会返回每个true值的完整坐标。例如对于2D张量，返回的每一行是[行索引, 列索引]，你可以直接遍历这些坐标对来处理：

const bool2DTensor = tf.tensor2d([[true, false], [false, true]]);

tf.tidy(() => {
  const coordsTensor = tf.where(bool2DTensor);
  coordsTensor.array().then(coordsArray => {
    coordsArray.forEach(([row, col]) => {
      console.log(`找到true值的坐标：行${row}，列${col}`);
    });
  });
});

性能优势 #

这个方法的核心优势在于：

• 所有索引提取逻辑都在TensorFlow的计算引擎中执行，不管是CPU还是GPU都能得到优化，避免了纯JS遍历的开销；
• 最终需要转成JS数组的只有索引数据，数据量远小于原始大张量，大幅降低了数据拷贝的耗时。

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