Merge branch 'master' into updates

brainpy/nn/nodes/ANN/batch_norm.py

@@ -1,11 +1,5 @@
 # -*- coding: utf-8 -*-
 
-"""
-adapted from jax.example_libraries.stax.BatchNorm
-https://jax.readthedocs.io/en/latest/_modules/jax/example_libraries/stax.html#BatchNorm
-"""
-
-
 from typing import Union
 
 import jax.nn
@@ -29,14 +23,23 @@ class BatchNorm(Node):
   Most commonly, the first axis of the data is the batch, and the last is
   the channel. However, users can specify the axes to be normalized.
 
+  adapted from jax.example_libraries.stax.BatchNorm
+  https://jax.readthedocs.io/en/latest/_modules/jax/example_libraries/stax.html#BatchNorm
+
   Parameters
   ----------
-  axis: axes where the data will be normalized. The axis of channels should be excluded.
-  epsilon: a value added to the denominator for numerical stability. Default: 1e-5
-  translate: whether to translate data in refactoring
-  scale: whether to scale data in refactoring
-  beta_init: an initializer generating the original translation matrix
-  gamma_init: an initializer generating the original scaling matrix
+  axis: int, tuple, list
+    axes where the data will be normalized. The axis of channels should be excluded.
+  epsilon: float
+    a value added to the denominator for numerical stability. Default: 1e-5
+  translate: bool
+    whether to translate data in refactoring
+  scale: bool
+    whether to scale data in refactoring
+  beta_init: brainpy.init.Initializer
+    an initializer generating the original translation matrix
+  gamma_init: brainpy.init.Initializer
+    an initializer generating the original scaling matrix
   """
   def __init__(self,
                axis: Union[int, tuple, list],
@@ -86,10 +89,14 @@ class BatchNorm1d(BatchNorm):
     axes where the data will be normalized. The axis of channels should be excluded.
   epsilon: float
     a value added to the denominator for numerical stability. Default: 1e-5
-  translate: whether to translate data in refactoring
-  scale: whether to scale data in refactoring
-  beta_init: an initializer generating the original translation matrix
-  gamma_init: an initializer generating the original scaling matrix
+  translate: bool
+    whether to translate data in refactoring
+  scale: bool
+    whether to scale data in refactoring
+  beta_init: brainpy.init.Initializer
+    an initializer generating the original translation matrix
+  gamma_init: brainpy.init.Initializer
+    an initializer generating the original scaling matrix
   """
   def __init__(self, axis=(0, 1), **kwargs):
     super(BatchNorm1d, self).__init__(axis=axis, **kwargs)
@@ -138,20 +145,24 @@ def _check_input_dim(self):
 
 class BatchNorm3d(BatchNorm):
   """3-D batch normalization.
-      The data should be of `(b, h, w, d, c)`, where `b` is the batch dimension,
-      `h` is the height dimension, `w` is the width dimension, `d` is the depth
-      dimension, and `c` is the channel dimension.
-
-      Parameters
-      ----------
-      axis: int, tuple, list
-        axes where the data will be normalized. The axis of channels should be excluded.
-      epsilon: float
-        a value added to the denominator for numerical stability. Default: 1e-5
-      translate: whether to translate data in refactoring
-      scale: whether to scale data in refactoring
-      beta_init: an initializer generating the original translation matrix
-      gamma_init: an initializer generating the original scaling matrix
+    The data should be of `(b, h, w, d, c)`, where `b` is the batch dimension,
+    `h` is the height dimension, `w` is the width dimension, `d` is the depth
+    dimension, and `c` is the channel dimension.
+
+    Parameters
+    ----------
+    axis: int, tuple, list
+      axes where the data will be normalized. The axis of channels should be excluded.
+    epsilon: float
+      a value added to the denominator for numerical stability. Default: 1e-5
+    translate: bool
+      whether to translate data in refactoring
+    scale: bool
+      whether to scale data in refactoring
+    beta_init: brainpy.init.Initializer
+      an initializer generating the original translation matrix
+    gamma_init: brainpy.init.Initializer
+      an initializer generating the original scaling matrix
       """
   def __init__(self, axis=(0, 1, 2, 3), **kwargs):
     super(BatchNorm3d, self).__init__(axis=axis, **kwargs)

brainpy/nn/nodes/ANN/tests/test_batchnorm.py

@@ -76,26 +76,26 @@ def test_batchnorm1(self):
 
 		print(model(inputs))
 
-	# def test_batchnorm2(self):
-	# 	i = bp.nn.Input((3, 4))
-	# 	b = bp.nn.BatchNorm(axis=(0, 2))  # channel axis: 1
-	# 	f = bp.nn.Reshape((-1, 12))
-	# 	o = bp.nn.GeneralDense(2)
-	# 	model = i >> b >>f >> o
-	# 	model.initialize(num_batch=2)
-	#
-	# 	inputs = bp.math.ones((2, 3, 4))
-	# 	inputs[0, 0, :] = 2.
-	# 	inputs[0, 1, 0] = 5.
-	# 	# print(inputs)
-	# 	print(model(inputs))
-	#
-	#
-	# 	X = bp.math.random.random((1000, 10, 3, 4))
-	# 	Y = bp.math.random.randint(0, 2, (1000, 10,  2))
-	# 	trainer = bp.nn.BPTT(model,
-	# 	                     loss=bp.losses.cross_entropy_loss,
-	# 	                     optimizer=bp.optim.Adam(lr=1e-3))
-	# 	trainer.fit([X, Y])
+	def test_batchnorm2(self):
+		i = bp.nn.Input((3, 4))
+		b = bp.nn.BatchNorm(axis=(0, 2))  # channel axis: 1
+		f = bp.nn.Reshape((-1, 12))
+		o = bp.nn.GeneralDense(2)
+		model = i >> b >>f >> o
+		model.initialize(num_batch=2)
+
+		inputs = bp.math.ones((2, 3, 4))
+		inputs[0, 0, :] = 2.
+		inputs[0, 1, 0] = 5.
+		# print(inputs)
+		print(model(inputs))
+
+
+		X = bp.math.random.random((1000, 10, 3, 4))
+		Y = bp.math.random.randint(0, 2, (1000, 10,  2))
+		trainer = bp.nn.BPTT(model,
+		                     loss=bp.losses.cross_entropy_loss,
+		                     optimizer=bp.optim.Adam(lr=1e-3))
+		trainer.fit([X, Y])
 
 

extensions/brainpylib/__init__.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 
-__version__ = "0.0.4"
+__version__ = "0.0.5"
 
 # IMPORTANT, must import first
 from . import register_custom_calls

extensions/changelog.rst

@@ -1,6 +1,17 @@
 Release notes (brainpylib)
 ##########################
 
+Version 0.0.5
+=============
+
+- Support operator customization on GPU by ``numba``
+
+
+Version 0.0.4
+=============
+
+- Support operator customization on CPU by ``numba``
+
 
 Version 0.0.3
 =============

extensions/setup.py

@@ -34,7 +34,7 @@
   author_email='[email protected]',
   packages=find_packages(exclude=['lib*']),
   include_package_data=True,
-  install_requires=["jax", "jaxlib", "pybind11>=2.6, <2.8"],
+  install_requires=["jax", "jaxlib", "pybind11>=2.6, <2.8", "cffi", "numba"],
   extras_require={"test": "pytest"},
   python_requires='>=3.7',
   url='https://github.com/PKU-NIP-Lab/BrainPy',

extensions/setup_cuda.py

@@ -90,7 +90,7 @@ def build_extension(self, ext):
   author_email='[email protected]',
   packages=find_packages(exclude=['lib*']),
   include_package_data=True,
-  install_requires=["jax", "jaxlib"],
+  install_requires=["jax", "jaxlib", "pybind11>=2.6, <2.8", "cffi", "numba"],
   extras_require={"test": "pytest"},
   python_requires='>=3.7',
   url='https://github.com/PKU-NIP-Lab/BrainPy',

extensions/setup_mac.py

@@ -22,7 +22,7 @@
                     sources=["lib/cpu_ops.cc"] + glob.glob("lib/*_cpu.cc"),
                     cxx_std=11,
                     # extra_link_args=["-rpath", "/Users/ztqakita/miniforge3/lib"], # m1
-                    extra_link_args=["-rpath", "/Users/ztqakita/miniforge3/lib"], # intel
+                    extra_link_args=["-rpath", "/Users/ztqakita/opt/miniconda3/lib"], # intel
                     define_macros=[('VERSION_INFO', __version__)]),
 ]
 
@@ -36,7 +36,7 @@
   author_email='[email protected]',
   packages=find_packages(exclude=['lib*']),
   include_package_data=True,
-  install_requires=["jax", "jaxlib", "pybind11>=2.6, <2.8"],
+  install_requires=["jax", "jaxlib", "pybind11>=2.6, <2.8", "cffi", "numba"],
   extras_require={"test": "pytest"},
   python_requires='>=3.7',
   url='https://github.com/PKU-NIP-Lab/BrainPy',