我的代码中存在明显的内存泄漏，导致使用的内存从5GB增加到15.7GB，仅在40到60秒内，然后程序因内存溢出错误而崩溃。我认为这种情况发生在创建数据集的张量时，而不是在训练模型时。我的数据包括本地存储的25,000张图片。因此，我使用了tensorflow.js内置的函数tf.data.generator(generator)来创建数据集，如这里所述。我认为这是创建大型数据集的最佳和最有效的方法，如这里所提到的。
示例
我使用了一个辅助类通过传入图片路径来创建我的数据集
class Dataset{ constructor(dirPath){ this.paths = this.#generatePaths(dirPath); } // 为所有要读取为缓冲区的图片生成文件路径 #generatePaths = (dirPath) => { const dir = fs.readdirSync(dirPath, {withFileTypes: true}) .filter(dirent => dirent.isDirectory()) .map(folder => folder.name) let imagePaths = []; dir.forEach(folder => { fs.readdirSync(path.join(dirPath, folder)).filter(file => { return path.extname(file).toLocaleLowerCase() === '.jpg' }).forEach(file => { imagePaths.push(path.resolve(path.join(dirPath, folder, file))) }) }) return imagePaths; } // 将图片缓冲区转换为Tensor对象 #generateTensor = (imagePath) => { const buffer = fs.readFileSync(imagePath); return tf.node.decodeJpeg(buffer, 3) .resizeNearestNeighbor([128, 128]) .toFloat() .div(tf.scalar(255.0)) } // 用相应的类别标记数据 #labelArray(index){return Array.from({length: 2}, (_, k) => k === index ? 1 : 0)}; // 传递给tf.data.generator()的Javascript生成器函数 * #imageGenerator(){ for(let i=0; i<this.paths.length; ++i){ let image; try { image = this.#generateTensor(this.paths[i]); } catch (error) { continue; } console.log(tf.memory()); yield image; } } // 传递给tf.data.generator()的Javascript生成器函数 * #labelGenerator(){ for(let i=0; i<this.paths.length; ++i){ const classIndex = (path.basename(path.dirname(this.paths[i])) === 'Cat' ? 0 : 1); const label = tf.tensor1d(this.#labelArray(classIndex), 'int32') console.log(tf.memory()); yield label; } } // 加载数据 loadData = () => { console.log('\n\n加载数据...') const xs = tf.data.generator(this.#imageGenerator.bind(this)); const ys = tf.data.generator(this.#labelGenerator.bind(this)); const ds = tf.data.zip({xs, ys}).batch(32).shuffle(32); return ds; }}我这样创建我的数据集:
const trainDS = new dataset(trainPath).loadData();问题
我知道tfjs中有内置的方法来管理内存，例如tf.tidy()和tf.dispose()。然而，我无法以阻止内存泄漏的方式实现它们，因为张量是由tf.data.generator函数生成的。
如何成功地在生成器yield张量后从内存中释放它们？
回答：
你创建的每个张量都需要释放 – 没有你习惯的JavaScript中的垃圾回收。因为张量不存储在JavaScript内存中（它们可能在GPU内存或WASM模块中等），所以JavaScript引擎无法跟踪它们。它们更像是指针而不是普通变量。
例如，在你的代码中:
return tf.node.decodeJpeg(buffer, 3) .resizeNearestNeighbor([128, 128]) .toFloat() .div(tf.scalar(255.0))每一步链式操作都会创建一个中间张量，这些张量从未被释放
这样读:
const decoded = tf.node.decodeJpeg(buffer, 3)const resized = decoded.resizeNearestNeighbor([128, 128])const casted = resized.toFloat();const normalized = casted.div(tf.scalar(255.0))return normalized;所以你分配了4个大型张量
你缺少的是
tf.dispose([decoded, resized, casted]);并且当你处理完图片后，也要tf.dispose(image)，它会释放normalized
对于所有张量都应如此操作。
我知道tfjs中有内置的方法来管理内存，例如tf.tidy()和tf.dispose()。然而，我无法以阻止内存泄漏的方式实现它们，因为张量是由tf.data.generator函数生成的。
你说你知道，但你实际上做的事情是创建了未释放的中间张量。
你可以通过将这样的函数包装在tf.tidy()中来帮助自己，它创建了一个本地作用域，因此所有未返回的内容都会自动释放。
例如:
#generateTensor = tf.tidy(imagePath) => { const buffer = fs.readFileSync(imagePath); return tf.node.decodeJpeg(buffer, 3) .resizeNearestNeighbor([128, 128]) .toFloat() .div(tf.scalar(255.0)) }这意味着中间张量会被释放，但你仍然需要在使用完返回值后释放它
