javascript - Tensorflow.js : How to manage memory for a large dataset created using tf. data.generator 内存泄漏？

Question

我的代码中存在明显的内存泄漏，导致我使用的内存在 40-60 秒内从 5GB 增加到 15.7GB，然后因 OOM 错误而使我的程序崩溃。我相信这种情况发生在我创建张量来形成数据集时，而不是在我训练模型时。我的数据由本地存储的 25,000 张图像组成。因此，我使用了here中描述的内置tensorflow.js函数tf.data.generator(generator)创建数据集。我相信这是创建大型数据集的最佳且最有效的方法，如上所述 here .

示例

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

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) 来处理标准化

对于张量的所有内容也是如此。

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() 中来帮助自己创建本地作用域，以便自动释放未返回的所有内容。

例如:

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

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