TensorFlow.js 中关于 tf.Model 的内存管理

Question

我是 TensorFlow 的新手。

https://js.tensorflow.org/tutorials/core-concepts.html中的"Memory Management: dispose and tf.tidy"部分说我们必须以特殊的方式管理内存。

然而，tfjs-layers中的类（例如tf.Model和Layer）似乎没有dispose和tf.tidy不接受这些作为返回值。

所以我的问题是：

tf.Model 会自动管理记忆吗？
如果没有，我该如何正确管理记忆？

示例代码：

function defineModel(
    regularizerRate: number,
    learningRate: number,
    stateSize: number,
    actionSize: number,
): tf.Model {
    return tf.tidy(() => { // Compile error here, I couldn't return model.
        const input = tf.input({
            name: "INPUT",
            shape: [stateSize],
            dtype: "int32" as any, // TODO(mysticatea): https://github.com/tensorflow/tfjs/issues/120
        })
        const temp = applyHiddenLayers(input, regularizerRate)
        const valueOutput = applyValueLayer(temp, regularizerRate)
        const policyOutput = applyPolicyLayer(temp, actionSize, regularizerRate)
        const model = tf.model({
            inputs: [input],
            outputs: [valueOutput, policyOutput],
        })

        // TODO(mysticatea): https://github.com/tensorflow/tfjs/issues/98
        model.compile({
            optimizer: tf.train.sgd(LEARNING_RATE),
            loss: ["meanSquaredError", "meanSquaredError"],
        })
        model.lossFunctions[1] = softmaxCrossEntropy

        return model
    })
}

Answer 1

您应该只在直接操作张量时使用 tf.tidy()。

构建模型时，您还没有直接操纵张量，而是在设置层如何组合在一起的结构。这意味着您不需要将模型创建包装在 tf.tidy().

中

只有当你调用"predict()"或"fit()"时，我们才处理具体的Tensor值，需要处理内存管理。

当调用"predict()"时，它return是一个Tensor，你必须处理它，或者用"tidy()"包围起来。

在 "fit()" 的情况下，我们在内部为您完成所有内存管理。 "fit()" 的 return 值是纯数字，因此您不需要将其包装在 "tidy()".

中

TensorFlow.js 中关于 tf.Model 的内存管理

Memory Management about tf.Model in TensorFlow.js

javascript

tensorflow

tensorflow.js