tensorflow.js 中计算直方图的最佳方法

Question

我需要使用 tensorflow.js 从一维大张量计算直方图。我找到了这样做的方法，但由于 oneHot 操作中的高内存消耗，它不是最佳的。我的代码示例如下：

 for (let index = 0; index < 50; index++) { //repeat some times to ensure no memory leaks
   const hist = getnormalHist(index);
   const histArray = (await hist.array()) as number[];
   const values = histArray.map((v, i) => ({ index: i, value: v }));
   await tfvis.render.barchart({ name: "hist" }, values);
 }

 function getnormalHist(seed: number) {
   return tf.tidy(() => {
     const rand = tf
       .randomNormal([100000], 0, 100, "int32", seed) //Generates long array of normal distributed randoms
       .add(500)
       .toInt();
     const oneHot = rand.oneHot(1000, 1, 0); // convert to oneHot makes it x1000 bigger
     const hist = oneHot.transpose().sum(1); // finally get histogram tensor
     return hist;
   });
 }

我需要使这段代码更快并减少内存消耗，但我不明白该怎么做。

Answer 1

tf.randomNormal([100000], 0, 100, "int32", seed) //Generates long array of normal distributed randoms
       .add(500)
       .toInt()

上述操作生成了一个样本分布，均值在 500 左右。可以使用均值 500 而不是 0 来生成它。添加操作不是必需的。此外，使用 toInt 是多余的，因为 dtype 已经是 int32。因此可以这样简化：

tf.randomNormal([100000], 500, 100, "int32", seed)

指定 0 和 1 onvalue 和 offvalue 是多余的，因为它们是默认值。

rand.oneHot(1000)

在计算每个索引的总和之前不需要转置。计算 0 轴上的总和将计算每个索引。不再使用大小为 100000 * 1000 的中间矩阵。

oneHot.sum(0)

作为总结，getNormalHist 将如下所示：

function getnormalHist(seed: number) {
   return tf.tidy(() => {
     const rand = tf
       .randomNormal([100000], 500, 100, "int32", seed) //Generates long array of normal distributed randoms

     const oneHot = rand.oneHot(1000); // convert to oneHot makes it x1000 bigger
     const hist = oneHot.sum(0); // finally get histogram tensor
     return hist;
   });
 }

Answer 2

我的解决方案是创建自定义 webGL 操作。

class HistProgram {
  variableNames = ["X"];
  outputShape: number[];
  userCode: string;

  constructor(numIndices: number, counts: number) {
    this.outputShape = [numIndices];
    this.userCode = `
      void main() {
        int ch = getOutputCoords();
        int c = 0;
        for (int i = 0; i < ${counts}; i++){
            if(ch == int(getX(i))){
                c++;
            }
        }
        setOutput(float(c));
      }
    `;
  }
}

现在我可以这样使用了:

const counts = 100000;
const channels = 1000;
const rand = tf.randomNormal([counts], 500, 100, "int32")
const prog = new HistProgram(channels, counts);
const hist = await tf.backend().compileAndRun(prog, [rand]).array();

Answer 3

现在我找到的最佳解决方案是：

const counts = 100000;
const channels = 1000;
const rand = tf.randomNormal([counts], 500, 100, "int32");
const hist = tf.sparseToDense(rand, tf.onesLike(rand), [channels], 0);

tensorflow.js 中计算直方图的最佳方法

Optimal way for calc histogram in tensorflow.js

tensorflow.js