add multiprocessing functions for batch running of BrainPy functions (#298)

add multiprocessing functions for batch running of BrainPy functions

Author: chaoming0625

brainpy/running/__init__.py

@@ -2,9 +2,26 @@
 
 
 """
-This module provides APIs for brain simulations.
+This module provides APIs for parallel brain simulations.
 """
 
-from .multiprocess import *
+from . import jax_multiprocessing
+from . import native_multiprocessing
+from . import pathos_multiprocessing
+from . import runner
+from . import constants
+
+
+__all__ = (native_multiprocessing.__all__ +
+           pathos_multiprocessing.__all__ +
+           jax_multiprocessing.__all__ +
+           runner.__all__ +
+           constants.__all__)
+
+
 from .runner import *
+from .jax_multiprocessing import *
+from .native_multiprocessing import *
+from .pathos_multiprocessing import *
 from .constants import *
+

brainpy/running/jax_multiprocessing.py

@@ -0,0 +1,141 @@
+# -*- coding: utf-8 -*-
+
+from typing import Sequence, Dict, Union
+
+import numpy as np
+from jax import vmap, pmap
+from jax.tree_util import tree_unflatten, tree_flatten
+
+import brainpy.math as bm
+from brainpy.types import Array
+
+__all__ = [
+  'jax_vectorize_map',
+  'jax_parallelize_map',
+]
+
+
+def jax_vectorize_map(
+    func: callable,
+    arguments: Union[Dict[str, Array], Sequence[Array]],
+    num_parallel: int,
+    clear_buffer: bool = False
+):
+  """Perform a vectorized map of a function by using ``jax.vmap``.
+
+  This function can be used in CPU or GPU backends. But it is highly
+  suitable to be used in GPU backends. This is because ``jax.vmap``
+  can parallelize the mapped axis on GPU devices.
+
+  Parameters
+  ----------
+  func: callable, function
+    The function to be mapped.
+  arguments: sequence, dict
+    The function arguments, used to define tasks.
+  num_parallel: int
+    The number of batch size.
+  clear_buffer: bool
+    Clear the buffer memory after running each batch data.
+
+  Returns
+  -------
+  results: Any
+    The running results.
+  """
+  if not isinstance(arguments, (dict, tuple, list)):
+    raise TypeError(f'"arguments" must be sequence or dict, but we got {type(arguments)}')
+  elements, tree = tree_flatten(arguments, is_leaf=lambda a: isinstance(a, bm.JaxArray))
+  if clear_buffer:
+    elements = [np.asarray(ele) for ele in elements]
+  num_pars = [len(ele) for ele in elements]
+  if len(np.unique(num_pars)) != 1:
+    raise ValueError(f'All elements in parameters should have the same length. '
+                     f'But we got {tree_unflatten(tree, num_pars)}')
+
+  res_tree = None
+  results = None
+  vmap_func = vmap(func)
+  for i in range(0, num_pars[0], num_parallel):
+    run_f = vmap(func) if clear_buffer else vmap_func
+    if isinstance(arguments, dict):
+      r = run_f(**tree_unflatten(tree, [ele[i: i + num_parallel] for ele in elements]))
+    else:
+      r = run_f(*tree_unflatten(tree, [ele[i: i + num_parallel] for ele in elements]))
+    res_values, res_tree = tree_flatten(r, is_leaf=lambda a: isinstance(a, bm.JaxArray))
+    if results is None:
+      results = tuple([np.asarray(val) if clear_buffer else val] for val in res_values)
+    else:
+      for j, val in enumerate(res_values):
+        results[j].append(np.asarray(val) if clear_buffer else val)
+    if clear_buffer:
+      bm.clear_buffer_memory()
+  if res_tree is None:
+    return None
+  results = ([np.concatenate(res, axis=0) for res in results]
+             if clear_buffer else
+             [bm.concatenate(res, axis=0) for res in results])
+  return tree_unflatten(res_tree, results)
+
+
+def jax_parallelize_map(
+    func: callable,
+    arguments: Union[Dict[str, Array], Sequence[Array]],
+    num_parallel: int,
+    clear_buffer: bool = False
+):
+  """Perform a parallelized map of a function by using ``jax.pmap``.
+
+  This function can be used in multi- CPU or GPU backends.
+  If you are using it in a single CPU, please set host device count
+  by ``brainpy.math.set_host_device_count(n)`` before.
+
+  Parameters
+  ----------
+  func: callable, function
+    The function to be mapped.
+  arguments: sequence, dict
+    The function arguments, used to define tasks.
+  num_parallel: int
+    The number of batch size.
+  clear_buffer: bool
+    Clear the buffer memory after running each batch data.
+
+  Returns
+  -------
+  results: Any
+    The running results.
+  """
+  if not isinstance(arguments, (dict, tuple, list)):
+    raise TypeError(f'"arguments" must be sequence or dict, but we got {type(arguments)}')
+  elements, tree = tree_flatten(arguments, is_leaf=lambda a: isinstance(a, bm.JaxArray))
+  if clear_buffer:
+    elements = [np.asarray(ele) for ele in elements]
+  num_pars = [len(ele) for ele in elements]
+  if len(np.unique(num_pars)) != 1:
+    raise ValueError(f'All elements in parameters should have the same length. '
+                     f'But we got {tree_unflatten(tree, num_pars)}')
+
+  res_tree = None
+  results = None
+  vmap_func = pmap(func)
+  for i in range(0, num_pars[0], num_parallel):
+    run_f = pmap(func) if clear_buffer else vmap_func
+    if isinstance(arguments, dict):
+      r = run_f(**tree_unflatten(tree, [ele[i: i + num_parallel] for ele in elements]))
+    else:
+      r = run_f(*tree_unflatten(tree, [ele[i: i + num_parallel] for ele in elements]))
+    res_values, res_tree = tree_flatten(r, is_leaf=lambda a: isinstance(a, bm.JaxArray))
+    if results is None:
+      results = tuple([np.asarray(val) if clear_buffer else val] for val in res_values)
+    else:
+      for j, val in enumerate(res_values):
+        results[j].append(np.asarray(val) if clear_buffer else val)
+    if clear_buffer:
+      bm.clear_buffer_memory()
+  if res_tree is None:
+    return None
+  results = ([np.concatenate(res, axis=0) for res in results]
+             if clear_buffer else
+             [bm.concatenate(res, axis=0) for res in results])
+  return tree_unflatten(res_tree, results)

brainpy/running/multiprocess.py renamed to brainpy/running/native_multiprocessing.py

@@ -1,17 +1,17 @@
 # -*- coding: utf-8 -*-
 
+from typing import Union, Sequence, Dict
 import multiprocessing
 
-
 __all__ = [
   'process_pool',
   'process_pool_lock',
-  'vectorize_map',
-  'parallelize_map',
 ]
 
 
-def process_pool(func, all_params, num_process):
+def process_pool(func: callable,
+                 all_params: Union[Sequence, Dict],
+                 num_process: int):
   """Run multiple models in multi-processes.
 
   .. Note::
@@ -47,7 +47,9 @@ def process_pool(func, all_params, num_process):
   return [r.get() for r in results]
 
 
-def process_pool_lock(func, all_params, nb_process):
+def process_pool_lock(func: callable,
+                      all_params: Union[Sequence, Dict],
+                      num_process: int):
   """Run multiple models in multi-processes with lock.
 
   Sometimes, you want to synchronize the processes. For example,
@@ -71,11 +73,11 @@ def some_func(..., lock, ...):
 
   Parameters
   ----------
-  func : callable
+  func: callable
       The function to run model.
   all_params : list, tuple, dict
       The parameters of the function arguments.
-  nb_process : int
+  num_process : int
       The number of the processes.
 
   Returns
@@ -84,7 +86,7 @@ def some_func(..., lock, ...):
       Process results.
   """
   print('{} jobs total.'.format(len(all_params)))
-  pool = multiprocessing.Pool(processes=nb_process)
+  pool = multiprocessing.Pool(processes=num_process)
   m = multiprocessing.Manager()
   lock = m.Lock()
   results = []
@@ -99,14 +101,3 @@ def some_func(..., lock, ...):
   pool.close()
   pool.join()
   return [r.get() for r in results]
-
-
-def vectorize_map(func, all_params, num_thread):
-  pass
-
-
-def parallelize_map(func, all_params, num_process):
-  pass
-
-
-