Question about clearing memory

[PikaPei]

Hi BrainPy team!

I'd like to use BrainPy to build my model and use particle swarm optimization (PSO, by PySwarms) to find the best parameters. However, the memory usage increased quickly when re-initializing a new model within each iteration, resulting in the program eventually crashing. I thought it was because I kept generating new models without destroying them properly, so I tried the bm.clear_buffer_memory(), but this approach didn't help.

Below is my pseudo code (you can see I create 100*100 models in my script). I can provide more details if needed. I would appreciate any advice. Thank you!

import brainpy as bp
import brainpy.math as bm

bm.set_environment(x64=True)

class Model(bp.NeuGroup):
    def __init__(self, size, var1_tau):
        super().__init__(size=size)

        self.var1_tau = var1_tau
        self.var1 = bm.Variable(bm.zeros(self.num))
        self.input = bm.Variable(bm.zeros(self.num))

        self.integral = bp.odeint(self.dvar1, method="exp_auto")

    def dvar1(self, var1, t):
        dvar1dt = -var1 / self.var1_tau + self.input
        return dvar1dt

    def update(self):
        t = bp.share["t"]
        dt = bp.share["dt"]

        self.var1.value = self.integral(self.var1, t, dt=dt)
        self.input[:] = 0.

    def run(self, duration, input_current):
        self.runner = bp.DSRunner(self, monitors=["var1"], inputs=[("input", input_current, "iter")])
        self.runner.run(duration)

    def loss(self):
        # output a value from the comparison of self.var1 and the target

def run_optimization_pso(n_particles, n_iterations):
    # pseudo code, it is managed by pyswarms in the actual implementation

    for _ in range(n_iterations):
        for _ in range(n_particles):
            # initialize a model with new parameters
            model = Model(size=1, var1_tau=...)
            model.run(10000)  # 10s simulation

            # compute loss
            model.loss()

            ##### clear buffer #####
            bm.clear_buffer_memory()
            ########################

if __name__ == "__main__":
    run_optimization_pso(n_particles=100, n_iterations=100)

[chaoming0625]

Thanks for the question. I will give you an example later.

[chaoming0625]

Thanks for the report.

The device's memory is occupied by three parts: 1. data or arrays; 2. compiled functions (including a lot of constants); 3. brainpy objects.

Therefore, for each of your single run function, I recommend using the following command to clear them all.

import brainpy.math as bm

def run_one_model(pars):
   model = ....
   model.run()
   l = model.loss()
   l = np.asarray(l)

   bm.clear_name_cache(ignore_warn=True)
   bm.clear_buffer_memory(array=True, compilation=True)
   return l

I hope this message is of some help.

[PikaPei]

Thanks for the suggestion!

However, it didn't work. Even a simple example eventually caused a crash. Would it be somewhat OS-dependent? I'm using Linux.

import brainpy as bp
import brainpy.math as bm

bm.set_environment(x64=True)

class Model(bp.NeuGroup):
    def __init__(self, size):
        super().__init__(size=size)

        self.var1 = bm.Variable(bm.zeros(self.num))

        self.integral = bp.odeint(self.dvar1, method="exp_auto")

    def dvar1(self, var1, t):
        dvar1dt = -var1
        return dvar1dt

    def update(self):
        t = bp.share["t"]
        dt = bp.share["dt"]

        self.var1.value = self.integral(self.var1, t, dt=dt)

    def run(self, duration):
        self.runner = bp.DSRunner(self, monitors=["var1"])
        self.runner.run(duration)

if __name__ == "__main__":
    for _ in range(100):
        for _ in range(100):
            model = Model(size=50)
            model.run(20000)

        # clear cache after 100 iterations
        bm.clear_name_cache(ignore_warn=True)
        bm.clear_buffer_memory(array=True, compilation=True)

[chaoming0625]

Thanks for the report. Please use brainpy.math.for_loop. DSRunner may have a memory leak. I have rewritten your code as:

import numpy as np

import brainpy as bp
import brainpy.math as bm

bm.set_environment(x64=True)

class Model(bp.NeuGroup):
  def __init__(self, size):
    super().__init__(size=size)

    self.var1 = bm.Variable(bm.zeros(self.num))

    self.integral = bp.odeint(self.dvar1, method="exp_auto")

  def dvar1(self, var1, t):
    dvar1dt = -var1
    return dvar1dt

  def update(self):
    t = bp.share["t"]
    dt = bp.share["dt"]
    self.var1.value = self.integral(self.var1, t, dt=dt)

def one_iter(n=10):
  for _ in range(n):
    def step(i):
      model.step_run(i)
      bp.clear_input(model)
      return model.var1.value

    model = Model(size=50)
    bm.for_loop(step, np.arange(int(20000 / bm.get_dt())), progress_bar=False)

  print(f'clear cache after {n} iterations')
  bm.clear_name_cache(ignore_warn=False)
  bm.clear_buffer_memory('cpu', array=True, compilation=True)

if __name__ == "__main__":
  for _ in range(100):
    one_iter()

[PikaPei]

This solution works well for now. But memory usage still goes up slowly, even after using bm.for_loop with bm.clear_name_cache and bm.clear_buffer_memory in every iteration. In my case, it can handle 100 iterations without crashing, which is quite good. However, I think if it goes beyond 200 or 300 iterations, the program will probably shut down. Still, I am very thankful for the current progress!

[chaoming0625]

I'll close this issue. If the error still occurs, please open it at any time.

[PikaPei]

I figured out the memory issue, and it's from my model implementation. May I pose a question?

The code shown below is causing the memory to increase, mainly because I use self.solution["Time"] = t * bm.get_dt() to store data. I also think my original implementation with DSRunner had the same problem, which is that I put it inside the class, although I haven't tested it. I'm wondering if this implementation is not a good way to do it, and what your recommendations would be? Thank you!

import brainpy as bp
import brainpy.math as bm
import numpy as np

bm.set_environment(x64=True)

class Model(bp.NeuGroup):
    def __init__(self):
        super().__init__(size=1)
        self.duration = 10000
        self.var1 = bm.Variable(bm.zeros(self.num))
        self.integral = bp.odeint(self.dvar1, method="exp_auto")

    def dvar1(self, var1, t):
        dvar1dt = -var1
        return dvar1dt

    def update(self):
        t = bp.share["t"]
        dt = bp.share["dt"]
        self.var1.value = self.integral(self.var1, t, dt=dt)

    def step(self, i):
        self.step_run(i)
        bp.clear_input(self)
        return self.var1.value

    def run_for_loop(self, duration=None, **kwargs):
        t = np.arange(int(self.duration / bm.get_dt()))
        sol = bm.for_loop(self.step, t, **kwargs)
        self.solution = {}
        self.solution["Time"] = t * bm.get_dt()  # key problem
        self.solution["var1"] = sol

if __name__ == "__main__":
    for _ in range(10000):
        for _ in range(100):
            model = Model()
            model.run_for_loop(progress_bar=False)

        bm.clear_name_cache(ignore_warn=True)
        bm.clear_buffer_memory(array=True, compilation=True)
        print("Cleared name cache and buffer memory.")