Tensorflow.js 中的内存泄漏：如何管理使用 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) 也会处理 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))
    }

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