我的 NewGeneration 方法使用 Tensorflow.js 的内存泄漏问题

Question

我正在使用 Tensorflow.js 开发一个 JavaScript 库，以通过浏览器中的 NeuroEvolution 训练模型。在我的主类中，我有一个 NewGeneration 方法，我认为它每次都会造成 160Tensors 的内存泄漏。我是不是哪里搞砸了？

我在 newGen 函数、pickOne 函数以及它调用的所有函数上添加了一个 tf.tidy() 包装，但没有任何效果。 如果您需要查看更多代码，请告诉我。

这是 newGeneration 的函数（保存的代理是失败的代理，代理是未来的代理）

newGen(){
    for(let i = 0; i < this.popSize; i++)
        this.agents.push(this.pickOne(this.savedAgents))

     for(let agent of this.savedAgents)
        tf.dispose(agent.brain.model)

     this.savedAgents = []
     print(tf.memory())
 }

这里是pickOne函数

pickOne(oldAgents){
    getTheChildsIndex() //Fake function but no code related to the problem here
    let child = oldAgents[index];
    let agent = new Agent() //Agent class holds a brain, in which is a tf.sequential model
                            //And a body, which has nothing to do with tf

    tf.tidy(()=>{
        agent.brain = child.brain.copy()
        agent.brain.mutate() //Mutate has proved to have another leak, but now it's fixed
    })
    return agent;
}

正如我被问到的那样，这是变异代码以防万一

mutate(mutationRate){
    tf.tidy(()=>{   
        let mutatedWeights
        const weights = this.model.getWeights()
        mutatedWeights = []
        for(let i = 0; i < weights.length; i++){
            let tensor = weights[i]
            let shape = tensor.shape
            let values = tensor.dataSync().slice()
            for(let j = 0; j < values.length; j++){
                if(random(1) < mutationRate)
                    values[j] = randomGaussian()
            }
            let newTensor = tf.tensor(values, shape)
            mutatedWeights.push(newTensor)
        }
        this.model.setWeights(mutatedWeights)
    })
    return TheBrain
}

还有复制功能

copy(){
    let modelCopy
    tf.tidy(()=>{
        modelCopy = new NeuralNetwork() // A helper class thatholds the model
        const weights  = this.model.getWeights()
        let clonedWeights = []
        for(let i = 0; i < weights.length; i++)
            clonedWeights.push(weights[i].clone())
        modelCopy.model.setWeights(clonedWeights)    
    })
    return modelCopy
}

我希望该函数不会造成泄漏，因此会处理它创建的所有张量。

Answer 1

我终于自己解决了。创建代理时，我不需要在 pickOne 方法中启动大脑。