Question about clearing memory
Closed this issue · 7 comments
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)
Thanks for the question. I will give you an example later.
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 lI hope this message is of some help.
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)
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()
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!
I'll close this issue. If the error still occurs, please open it at any time.
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.")