尽管使用了适当的张量处理 (tfjs),但仍无法找到内存泄漏

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

我已经尝试了各种方法来处理张量(tf.dispose(),start/endscope)。 我最接近的是通过这段代码,每次执行后都会留下 1 个未使用的张量。该程序需要大约 2 小时 运行 足以用完 64 GB RAM(大内存泄漏)。

我还怀疑除了基于 TFJS 的操作之外还有其他因素导致内存泄漏,尽管(理论上)垃圾收集应该清理它。

下面的这段代码是一个由事件侦听器处理程序处理的事件。如果您对此问题有任何帮助,我们将不胜感激!

'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);

  },
};

首先,将模型加载和推理分开 - 在您当前的代码中,每次需要对新图像进行 运行 预测时都会重新加载模型。

然后查看预测函数中是否存在任何可能的漏洞 - 所以一旦模型加载完毕。

您正在加载模型并处理每一层,但这并不意味着模型本身被卸载,因此模型的一部分很有可能保留在内存中。

但泄漏本身就是这一行:

rawArray = await rawArray.data();

该变量已被使用并且它是一个张量。
现在你用数据数组覆盖同一个变量,张量永远不会被释放。