简体   繁体   中英

Memory Leak Issue with my NewGeneration method using Tensorflow.js

I am developing a library in JavaScript with Tensorflow.js to train models via NeuroEvolution in the browser. In my main class I have a NewGeneration method, and I think it creates a memory leak of 160Tensors every time. Did I mess up somewhere?

I have added a tf.tidy() wrapping on the newGen function, pickOne function, and all the ones it calls, but nothing works. If you need to see more code, please tell me.

Here's the function for the newGeneration (saved agents are ones that failed and agents are the future ones)

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())
 }

And here's the pickOne function

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;
}

As I have been asked, here's the mutate code just in case

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
}

As well as the copy function

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
}

I expect the function to not create the leak, hence disposing of all the Tensors it creates.

I've finally sorted it out by myself. I needed to NOT initiate the brain in the pickOne method when I created an agent.

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM