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

Question

我是 TensorFlow 的新手。

https://js.tensorflow.org/tutorials/core-concepts.html 中的“内存管理:dispose 和 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()”被调用时，它返回一个张量，您必须处理它，或者用“tidy()”包围它。

对于“fit()”，我们在内部为您完成所有内存管理。 “fit()”的返回值是纯数字，因此您不需要将其包装在“tidy()”中。