Skip to content

Parallel simulation of multiple neuron groups inside brainpy.DSRunner #490

@CloudyDory

Description

@CloudyDory

I have written a biological neural network that contains several groups of neurons and synapses in Brainpy. However, during simulation, it seems that the code cannot fully utilize all CPU cores. The code looks like this:

import brainpy as bp
import brainpy.math as bm

class Network(bp.DynSysGroup):
    def __init__(self):
        super().__init__()
        self.neu1 = bp.neurons.HH(size=10, **param1)
        self.neu2 = bp.neurons.HH(size=10, **param2)
        self.neu3 = bp.neurons.HH(size=10, **param3)
        ......
        self.neuron = [self.neu1, self.neu2, self.neu3, ...]

        self.syn1 = bp.synapses.AMPA(self.neu1, self.neu2, ...)
        self.syn2 = bp.synapses.AMPA(self.neu1, self.neu3, ...)
        self.syn3 = bp.synapses.GABAa(self.neu2, self.neu3, ...)
        ......
        self.synapse = [self.syn1, self.syn2, self.syn3, ...]

    def update(self, inputs):
        for neu in self.neuron:
            neu.update(inputs)
        
        for syn in self.synapse:
            syn.update()

if __name__ == '__main__':
    dt = 0.01  # ms
    bm.set_dt(dt)
            
    network = Network()
    I_input = bp.inputs.section_input(...)
    
    runner = bp.DSRunner(
        network, 
        monitors = ['neu1.V', 'neu2.V', ...], 
        jit = True,
        dt = dt,
        progress_bar = True
    )
    runner.run(inputs=I_input) 

I think updating each neuron group and synapse group in python for-loop is slow, In theory, the neuron and synapse groups can be run in parallel. I have tried modified the update function as this using the Python threading module:

def update(self, inputs):
        threads = []
        for neu in self.neuron:
            threads.append(threading.Thread(target=neu.update, args=(inputs,)))
            threads[-1].start()
        for thread in threads:
            thread.join()
        
        threads = []
        for syn in self.synapse:
            threads.append(threading.Thread(target=syn.update))
            threads[-1].start()
        for thread in threads:
            thread.join()

or this:

def update_parallel(self, neuron):
    neuron.update()

def update(self, inputs):
    bp.running.cpu_unordered_parallel(self.update_parallel, [[neu] for neu in self.neuron], num_process=min(len(self.neuron),16))
    
    for syn in self.synapse:
        syn.update()

But from the task manager (in Ubuntu 22.04), I find that the process is still running on a single core. Is there a way to run multiple neuron groups in a network in parallel in Brainpy?

Thanks!

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions