Tensorflow.js 中的内存泄漏：如何管理使用 tf.data.generator 创建的大型数据集的内存？

Question

我的代码中存在明显的内存泄漏，导致我使用的内存在 40-60 秒内从 5gb 增加到 15.7gb，然后我的程序因 OOM 错误而崩溃。我相信这发生在我创建张量以形成数据集时，而不是在我训练模型时。我的数据包含本地存储的 25,000 张图像。因此，我使用了 tensorflow.js 函数 tf.data.generator(generator) 描述的 here to create the dataset. I believe this is the best and most efficient way to create a large dataset as mentioned .

例子

我使用助手 class 通过传入图像路径来创建我的数据集

class Dataset{

    constructor(dirPath){
        this.paths = this.#generatePaths(dirPath);
    }

    // Generate file paths for all images to be read as buffer
    #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;
    }

    // Convert image buffer to a Tensor object
    #generateTensor = (imagePath) => {
        const buffer = fs.readFileSync(imagePath);
        return tf.node.decodeJpeg(buffer, 3)
            .resizeNearestNeighbor([128, 128])
            .toFloat()
            .div(tf.scalar(255.0))
    }

    // Label the data with the corresponding class
    #labelArray(index){return Array.from({length: 2}, (_, k) => k === index ? 1 : 0)};

    // Javascript generator function passed to tf.data.generator()
    * #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;
        }
    }

    // Javascript generator function passed to tf.data.generator()
    * #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;
        }
    }

    // Load data
    loadData = () => {
        console.log('\n\nLoading data...')
        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 函数生成的。

在生成器生成张量后，我将如何从内存中成功处理张量？

Answer 1

你创建的每个张量，你都需要处理 - 没有你在 JS 中习惯的垃圾收集。那是因为张量没有保存在 JS 内存中（它们可以在 GPU 内存或 WASM 模块等中），所以 JS 引擎无法跟踪它们。它们更像是指针而不是普通变量。

例如，在您的代码中：

        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

对于所有张量也是如此。

I am aware of built-in tfjs methods to manage memory such as tf.tidy() and tf.dispose(). However, I was unable to implement them in such a way to stop the memory leak, as the tensors are generated by the tf.data.generator function.

你说你知道，但你通过创建你没有处理的临时张量来做完全相同的事情。

您可以通过将此类函数包装在创建本地作用域的 tf.tidy() 中来帮助自己，这样所有未 returned 的内容都会自动释放。

例如：

   #generateTensor = tf.tidy(imagePath) => {
        const buffer = fs.readFileSync(imagePath);
        return tf.node.decodeJpeg(buffer, 3)
            .resizeNearestNeighbor([128, 128])
            .toFloat()
            .div(tf.scalar(255.0))
    }

这意味着临时张量将被处理掉，但您仍然需要在完成后处理 return 值

Tensorflow.js 中的内存泄漏：如何管理使用 tf.data.generator 创建的大型数据集的内存？

Memory leak in Tensorflow.js: How to manage memory for a large dataset created using tf.data.generator?

javascript

node.js

machine-learning

memory-leaks

tensorflow.js