Unable to find memory leak despite using proper tensor disposal (tfjs)

Question

I've tried all sorts of ways to dispose of tensors (tf.dispose(), start/endscope). The closest I've got was through this code, where 1 unused tensor remains after each execution. It takes about 2 hours for this program to run enough to use up 64 GB of RAM (big memory leak).

I also suspect that other factors besides TFJS-based operations are contributing to the memory leak, though (in theory) garbage collection should clean this up.

The piece of code below is one event that gets processed by an Event Listener handler. Any help with this issue would be greatly appreciated!

'use strict';

global.fetch = require("node-fetch");
const { MessageActionRow, MessageButton, Permissions } = require('discord.js');
const { mod, eco, m, n } = require(`../../index.js`);
const { Readable } = require('stream');
const PImage = require('pureimage');
const tf = require('@tensorflow/tfjs');
const tfnode = require('@tensorflow/tfjs-node');
const wait = require('util').promisify(setTimeout);

let bufferToStream = (binary) => {
  let readableInstanceStream = new Readable({
    read() {
      this.push(binary);
      this.push(null);
    }
  });
  return readableInstanceStream;
}

const predict = async (imageUrl, modelFile) => {

  let model = await tf.loadLayersModel(modelFile);
  let modelClasses = [ "NSFW", "SFW" ];

  let data = await fetch(imageUrl);
  let fileType = data.headers.get("Content-Type");
  let buffer = await data.buffer();

  let stream = bufferToStream(buffer);
  let image;
  if ((/png/).test(fileType)) {
    image = await PImage.decodePNGFromStream(stream);
  }
  else if ((/jpe?g/).test(fileType)) {
    image = await PImage.decodeJPEGFromStream(stream);
  }
  else {
    return;
  }

  let rawArray;
  rawArray = tf.tidy(() => {
    let tensorImage;
    tensorImage = tf.browser.fromPixels(image).toFloat();
    tensorImage = tf.image.resizeNearestNeighbor(tensorImage, [model.inputs[0].shape[1], model.inputs[0].shape[2]]);
    tensorImage = tensorImage.reshape([1, model.inputs[0].shape[1], model.inputs[0].shape[2], model.inputs[0].shape[3]]);

    return model.predict(tensorImage);
  });

  rawArray = await rawArray.data();
  rawArray = Array.from(rawArray);

  tf.disposeVariables();
  model.layers.forEach(l => {
    l.dispose();
  });

  if (rawArray[1] > rawArray[0]) {
    return [`SFW`, rawArray[1]];
  }
  else {
    return [`NSFW`, rawArray[0]];
  }
};

const getResults = async (imageLink, imageNumber) => {
  let image = `${imageLink}`;
  let prediction = await predict(image, `file://D:/retake7/sfwmodel/model.json`);
  let className = `SFW`;
  if (prediction[0] == `NSFW`) {
    className = `**NSFW**`;
  }
  return [`[Image ${imageNumber+1}](${imageLink}): ${className} (${(prediction[1]*100).toFixed(2)}% Certainty)`, ((prediction[1]*100).toFixed(2))*1];
}

const main = async (message, client, Discord) => {
  if (message.attachments.size == 0 || message.author.bot || message.channel.nsfw) return;

  await client.shard.broadcastEval(c => {
    console.log(`Scanning...`);
  }).catch(e => {
    return;
  });

  let inChannel = await eco.seid.get(`${message.guild.id}.${message.channel.id}.active`);
  let sfwImage = await eco.seid.get(`${message.guild.id}.sfwAlerts`);
  if (inChannel == `no`) return;

  let atmentArr = Array.from(message.attachments);
  let msgArr = [];
  if (message.attachments.size > 1) {
    msgArr.push(`**Images Scanned**`);
  } else {
    msgArr.push(`**Image Scanned**`);
  }
  let hasNSFW = false;
  let uncertain = false;

  for (i = 0; i < message.attachments.size; i++) {
    let msg = await getResults(atmentArr[i][1][`proxyURL`], i);
    if (msg[1] < 80) {
      uncertain = true;
    }
    if (msg[0].includes(`NSFW`)) {
      hasNSFW = true;
    }
    msgArr.push(msg[0]);
  }

  if (uncertain == false && hasNSFW == false) {
    let cont = `${msgArr.join(`\n`)}`;
    msgArr = null;
    client.seid.set(`${message.channel.id}.previousScan`, cont);
    return;
  }

  let embed = new Discord.MessageEmbed()
    .setColor(`GREEN`)
    .setDescription(msgArr.join(`\n`));

  let cont2 = `${msgArr.join(`\n`)}`;
  client.seid.set(`${message.channel.id}.previousScan`, cont2);
  msgArr = null;

  if (sfwImage != `no` || hasNSFW || msg[1] <= 80) {
    embed.setColor(`RED`);
    await message.delete();
    let msgSent = await message.channel.send({embeds: [embed], components: [row]});
  };
};

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

    await main(message, client, Discord);

  },
};

Answer 1

first, separate model loading and inference - in your current code, you'd reload a model each time you need to run prediction on a new image.

and then look at any possible leaks in prediction function - so once model is loaded.

you're loading a model and disposing each layer, but that doesn't mean model itself gets unloaded so there more than a chance that part of model remains in memory.

but leak itself is this line:

rawArray = await rawArray.data();

that variable is already used and its a tensor.
now you're overwriting the same variable with a data array and tensor never gets disposed.