tfjs-node memory leak even after proper tensor disposal

Question

I've struggling to find where a memory leak occurs in this file. This file is exported as an Event Listener. For context, I have 92 shards (meaning 92 of these listeners) running. I import the model from outside of this file so it's only loaded once per shard occurrence (stable 75 tensors in memory). However, after a few minutes, all the RAM on my computer is consumed (the function inside the file is called a dozen or so times per second). Have I overlooked any place which may cause this memory leak?

const use = require(`@tensorflow-models/universal-sentence-encoder`);
const tf = require(`@tensorflow/tfjs-node`);

const run = async (input, model) => {

    const useObj = await use.load();
    const encodings = [ await useObj.tokenizer.encode(input) ];

    const indicesArr = encodings.map(function (arr, i) { return arr.map(function (d, index) { return [i, index]; }); });
    var flattenedIndicesArr = [];
    for (i = 0; i < indicesArr.length; i++) {
      flattenedIndicesArr = flattenedIndicesArr.concat(indicesArr[i]);
    }

    const indices = tf.tensor2d(flattenedIndicesArr, [flattenedIndicesArr.length, 2], 'int32')
    const value = tf.tensor1d(tf.util.flatten([ encodings ]), 'int32')

    const prediction = await model.executeAsync({ Placeholder_1: indices, Placeholder: value });
    const classes = [ 'Identity Attack', 'Insult', 'Obscene', 'Severe Toxicity', 'Sexual Explicit', 'Threat', 'Toxicity' ]
    let finArr = [];
    let finMsg = `Input: ${input}, `;

    for (i = 0; i < prediction.length; i++) {
      const sorted = tf.topk(prediction[i], 2);
      const predictions = [ sorted.values.arraySync(), sorted.indices.arraySync() ];

      const percentage = (predictions[0][0][0]*100).toFixed(2);
      if (predictions[1][0][0] == 1) {
        finArr.push(`${classes[i]} (${percentage}%)`);
      }
      tf.dispose([ sorted, predictions ]);
    }
    for (i = 0; i < finArr.length; i++) {
      finMsg+=`${finArr[i]}, `;
    }

    tf.dispose([ prediction, indices, value, useObj ]);

    console.log(finMsg);
    console.log(tf.memory());
};

const main = async (message, client, Discord, model) => {
  if (message.author.bot) return;

  const input = message.content;
  await run(input, model);

};

module.exports = {
  event: 'messageCreate',
  run: async (message, client, Discord, model) => {

    await main(message, client, Discord, model);

  },
};

Answer 1

to start with, you say this runs multiple times - so why are you loading model again and again? and disposing model is tricky, big chance that's part of your memory leak.

move const useObj = await use.load() outside of run loop and don't dispose it until you're done with all of the runs.