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

mys*_*tea 3 javascript tensorflow tensorflow.js

我是TensorFlow的新手.

https://js.tensorflow.org/tutorials/core-concepts.html中的"内存管理:disposetf.tidy"部分说我们必须以特殊方式管理记忆.

但是,tfjs-layers(例如tf.ModelLayer)中的类似乎没有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
    })
}
Run Code Online (Sandbox Code Playgroud)

Nik*_*hil 6

你应该只在直接操纵张量时使用tf.tidy().

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

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

当调用"predict()"时,它返回一个Tensor,你必须处理它,或者用"tidy()"包围.

在"fit()"的情况下,我们在内部为您完成所有内存管理."fit()"的返回值是普通数字,因此您不需要将其包装在"tidy()"中.