diff --git a/docs/core_concept/brainpy_transform_concept-old.ipynb b/docs/core_concept/brainpy_transform_concept-old.ipynb deleted file mode 100644 index c8b3a771b..000000000 --- a/docs/core_concept/brainpy_transform_concept-old.ipynb +++ /dev/null @@ -1,654 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": { - "collapsed": true, - "jupyter": { - "outputs_hidden": true - } - }, - "source": [ - "# Concept 1: Object-oriented Transformation" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "@[Chaoming Wang](https://github.com/chaoming0625)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Most computation in BrainPy relies on [JAX](https://jax.readthedocs.io/en/latest/).\n", - "JAX has provided wonderful transformations, including differentiation, vecterization, parallelization and just-in-time compilation, for Python programs. If you are not familiar with it, please see its [documentation](https://jax.readthedocs.io/en/latest/).\n", - "\n", - "However, JAX only supports functional programming, i.e., transformations for Python functions. This is not what we want. Brain Dynamics Modeling need object-oriented programming.\n", - "\n", - "To meet this requirement, BrainPy defines the interface for object-oriented (OO) transformations. These OO transformations can be easily performed for BrainPy objects.\n", - "\n", - "In this section, let's talk about the BrainPy concept of object-oriented transformations." - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "metadata": {}, - "outputs": [], - "source": [ - "import brainpy as bp\n", - "import brainpy.math as bm\n", - "\n", - "# bm.set_platform('cpu')" - ] - }, - { - "cell_type": "code", - "execution_count": 5, - "metadata": { - "collapsed": false, - "jupyter": { - "outputs_hidden": false - } - }, - "outputs": [ - { - "data": { - "text/plain": [ - "'2.3.0'" - ] - }, - "execution_count": 5, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "bp.__version__" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Illustrating example: Training a network" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To illustrate this concept, we need a demonstration example. Here, we choose the popular neural network training as the illustrating case.\n", - "\n", - "In this training case, we want to teach the neural network to correctly classify a random array as two labels (`True` or `False`). That is, we have the training data:" - ] - }, - { - "cell_type": "code", - "execution_count": 6, - "metadata": {}, - "outputs": [], - "source": [ - "num_in = 100\n", - "num_sample = 256\n", - "X = bm.random.rand(num_sample, num_in)\n", - "Y = (bm.random.rand(num_sample) < 0.5).astype(float)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We use a two-layer feedforward network:" - ] - }, - { - "cell_type": "code", - "execution_count": 7, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Sequential(\n", - " [0] Linear0\n", - " [1] relu\n", - " [2] Linear1\n", - ")\n" - ] - } - ], - "source": [ - "class Linear(bp.BrainPyObject):\n", - " def __init__(self, n_in, n_out):\n", - " super().__init__()\n", - " self.num_in = n_in\n", - " self.num_out = n_out\n", - " init = bp.init.XavierNormal()\n", - " self.W = bm.Variable(init((n_in, n_out)))\n", - " self.b = bm.Variable(bm.zeros((1, n_out)))\n", - "\n", - " def __call__(self, x):\n", - " return x @ self.W + self.b\n", - "\n", - "\n", - "net = bp.Sequential(Linear(num_in, 20),\n", - " bm.relu,\n", - " Linear(20, 2))\n", - "print(net)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Here, we use a supervised learning training paradigm. " - ] - }, - { - "cell_type": "code", - "execution_count": 8, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Train 400 epoch, loss = 0.6710\n", - "Train 800 epoch, loss = 0.5992\n", - "Train 1200 epoch, loss = 0.5332\n", - "Train 1600 epoch, loss = 0.4720\n", - "Train 2000 epoch, loss = 0.4189\n", - "Train 2400 epoch, loss = 0.3736\n", - "Train 2800 epoch, loss = 0.3335\n", - "Train 3200 epoch, loss = 0.2972\n", - "Train 3600 epoch, loss = 0.2644\n", - "Train 4000 epoch, loss = 0.2346\n" - ] - } - ], - "source": [ - "rng = bm.random.RandomState(123)\n", - "\n", - "\n", - "# Loss function\n", - "@bm.to_object(child_objs=net, dyn_vars=rng)\n", - "def loss():\n", - " # shuffle the data\n", - " key = rng.split_key()\n", - " x_data = rng.permutation(X, key=key)\n", - " y_data = rng.permutation(Y, key=key)\n", - " # prediction\n", - " predictions = net(dict(), x_data)\n", - " # loss\n", - " l = bp.losses.cross_entropy_loss(predictions, y_data)\n", - " return l\n", - "\n", - "\n", - "# Gradient function\n", - "grad = bm.grad(loss, grad_vars=net.vars(), return_value=True)\n", - "\n", - "# Optimizer\n", - "optimizer = bp.optim.SGD(lr=1e-2, train_vars=net.vars())\n", - "\n", - "\n", - "# Training step\n", - "@bm.to_object(child_objs=(grad, optimizer))\n", - "def train(i):\n", - " grads, l = grad()\n", - " optimizer.update(grads)\n", - " return l\n", - "\n", - "\n", - "num_step = 400\n", - "for i in range(0, 4000, num_step):\n", - " # train 400 steps once\n", - " ls = bm.for_loop(train, operands=bm.arange(i, i + num_step))\n", - " print(f'Train {i + num_step} epoch, loss = {bm.mean(ls):.4f}')" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In the above example, we have seen classical elements in a neural network training, such as \n", - "\n", - "- `net`: neural network\n", - "- `loss`: loss function\n", - "- `grad`: gradient function\n", - "- `optimizer`: parameter optimizer\n", - "- `train`: training step" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In BrainPy, all these elements can be defined as class objects and can be used for performing OO transformations. " - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In essence, the concept of BrainPy object-oriented transformation has three components:\n", - "\n", - "- `BrainPyObject`: the base class for object-oriented programming\n", - "- `Variable`: the varibles in the class object, whose values are ready to be changed/updated during transformation\n", - "- `ObjectTransform`: the transformations for computation involving `BrainPyObject` and `Variable`" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## ``BrainPyObject`` and its ``Variable``" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "``BrainPyObject`` is the base class for object-oriented programming in BrainPy. \n", - "It can be viewed as a container which contains all needed [Variable](../tutorial_math/arrays_and_variables.ipynb) for our computation." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](./imgs/net_with_two_linear.png)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In the above example, ``Linear`` object has two ``Variable``: *W* and *b*. The ``net`` we defined is further composed of two ``Linear`` objects. We can expect that four variables can be retrieved from it." - ] - }, - { - "cell_type": "code", - "execution_count": 9, - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "dict_keys(['Linear0.W', 'Linear0.b', 'Linear1.W', 'Linear1.b'])" - ] - }, - "execution_count": 9, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "net.vars().keys()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "An important question is, **how to define `Variable` in a `BrainPyObject` so that we can retrieve all of them?**" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Actually, all Variable instance which can be accessed by `self.` attribue can be retrived from a `BrainPyObject` recursively. \n", - "No matter how deep the composition of ``BrainPyObject``, once `BrainPyObject` instance and their `Variable` instances can be accessed by `self.` operation, all of them will be retrieved. " - ] - }, - { - "cell_type": "code", - "execution_count": 10, - "metadata": {}, - "outputs": [], - "source": [ - "class SuperLinear(bp.BrainPyObject):\n", - " def __init__(self, ):\n", - " super().__init__()\n", - " self.l1 = Linear(10, 20)\n", - " self.v1 = bm.Variable(3)\n", - " \n", - "sl = SuperLinear()" - ] - }, - { - "cell_type": "code", - "execution_count": 11, - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "dict_keys(['SuperLinear0.v1', 'Linear2.W', 'Linear2.b'])" - ] - }, - "execution_count": 11, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "# retrieve Variable\n", - "sl.vars().keys()" - ] - }, - { - "cell_type": "code", - "execution_count": 12, - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "dict_keys(['SuperLinear0', 'Linear2'])" - ] - }, - "execution_count": 12, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "# retrieve BrainPyObject\n", - "sl.nodes().keys()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "However, we cannot access the ``BrainPyObject`` or ``Variable`` which is in a Python container (like tuple, list, or dict). For this case, we can register our objects and variables through ``.register_implicit_vars()`` and ``.register_implicit_nodes()``:" - ] - }, - { - "cell_type": "code", - "execution_count": 13, - "metadata": {}, - "outputs": [], - "source": [ - "class SuperSuperLinear(bp.BrainPyObject):\n", - " def __init__(self, register=False):\n", - " super().__init__()\n", - " self.ss = [SuperLinear(), SuperLinear()]\n", - " self.vv = {'v_a': bm.Variable(3)}\n", - " if register:\n", - " self.register_implicit_nodes(self.ss)\n", - " self.register_implicit_vars(self.vv)" - ] - }, - { - "cell_type": "code", - "execution_count": 14, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "dict_keys([])\n", - "dict_keys(['SuperSuperLinear0'])\n" - ] - } - ], - "source": [ - "# without register\n", - "ssl = SuperSuperLinear(register=False)\n", - "print(ssl.vars().keys())\n", - "print(ssl.nodes().keys())" - ] - }, - { - "cell_type": "code", - "execution_count": 15, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "dict_keys(['SuperSuperLinear1.v_a', 'SuperLinear3.v1', 'SuperLinear4.v1', 'Linear5.W', 'Linear5.b', 'Linear6.W', 'Linear6.b'])\n", - "dict_keys(['SuperSuperLinear1', 'SuperLinear3', 'SuperLinear4', 'Linear5', 'Linear6'])\n" - ] - } - ], - "source": [ - "# with register\n", - "ssl = SuperSuperLinear(register=True)\n", - "print(ssl.vars().keys())\n", - "print(ssl.nodes().keys())" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Transform a function to `BrainPyObject`" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](./imgs/loss_with_net_and_rng.png)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Let's go back to our network training.\n", - "After the definition of `net`, we further define a ``loss`` function whose computation involves the ``net`` object for neural network prediction and a ``rng`` Variable for data shuffling. \n", - "\n", - "This Python function is then transformed into a ``BrainPyObject`` instance by ``brainpy.math.to_object`` interface. " - ] - }, - { - "cell_type": "code", - "execution_count": 16, - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "FunAsObject(nodes=[Sequential0],\n", - " num_of_vars=1)" - ] - }, - "execution_count": 16, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "loss" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "All `Variable` used in this instance can also be retrieved through:" - ] - }, - { - "cell_type": "code", - "execution_count": 17, - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "dict_keys(['loss0._var0', 'Linear0.W', 'Linear0.b', 'Linear1.W', 'Linear1.b'])" - ] - }, - "execution_count": 17, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "loss.vars().keys()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Note that, when using `to_object()`, we need to explicitly declare all `BrainPyObject` and `Variable` used in this Python function. \n", - "Due to the recursive retrieval property of `BrainPyObject`, we only need to specify the latest composition object.\n", - "\n", - "In the above `loss` object, we do not need to specify two ``Linear`` object. Instead, we only need to give the top level object ``net`` into ``to_object()`` transform. \n", - "\n", - "Similarly, when we transform ``train`` function into a ``BrainPyObject``, we just need to point out the ``grad`` and ``opt`` we have used, rather than the previous *loss*, *net* or *rng*. " - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](./imgs/train_with_grad_and_opt.png)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## BrainPy object-oriented transformations" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "BrainPy object-oriented transformations are designed to work on ``BrainPyObject``. \n", - "These transformations include autograd ``brainpy.math.grad()`` and JIT ``brainpy.math.jit()``." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In our case, we used two OO transformations provided in BrainPy. \n", - "\n", - "First, ``grad`` object is defined with the ``loss`` function. Within it, we need to specify what variables we need to compute their gradients through `grad_vars`. " - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Note that, the OO transformation of any ``BrainPyObject`` results in another ``BrainPyObject`` object. Therefore, it can be recersively used as a component to form the larger scope of object-oriented programming and object-oriented transformation. " - ] - }, - { - "cell_type": "code", - "execution_count": 18, - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "GradientTransform(target=loss0, \n", - " num_of_grad_vars=4, \n", - " num_of_dyn_vars=1)" - ] - }, - "execution_count": 18, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "grad" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](./imgs/grad_with_loss.png)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Next, we train 400 steps once by using a ``for_loop`` transformation. Different from ``grad`` which return a `BrainPyObject` instance, `for_loop` direactly returns the loop results. " - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](./imgs/for-loop-train.png)" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "BrainPy", - "language": "python", - "name": "brainpy" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.6" - }, - "latex_envs": { - "LaTeX_envs_menu_present": true, - "autoclose": false, - "autocomplete": true, - "bibliofile": "biblio.bib", - "cite_by": "apalike", - "current_citInitial": 1, - "eqLabelWithNumbers": true, - "eqNumInitial": 1, - "hotkeys": { - "equation": "Ctrl-E", - "itemize": "Ctrl-I" - }, - "labels_anchors": false, - "latex_user_defs": false, - "report_style_numbering": false, - "user_envs_cfg": false - }, - "toc": { - "base_numbering": 1, - "nav_menu": {}, - "number_sections": false, - "sideBar": true, - "skip_h1_title": false, - "title_cell": "Table of Contents", - "title_sidebar": "Contents", - "toc_cell": false, - "toc_position": { - "height": "calc(100% - 180px)", - "left": "10px", - "top": "150px", - "width": "245.75px" - }, - "toc_section_display": true, - "toc_window_display": true - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/core_concept/imgs/for-loop-train.png b/docs/core_concept/imgs/for-loop-train.png deleted file mode 100644 index 5c380e5a7..000000000 Binary files a/docs/core_concept/imgs/for-loop-train.png and /dev/null differ diff --git a/docs/core_concept/imgs/grad_with_loss.png b/docs/core_concept/imgs/grad_with_loss.png deleted file mode 100644 index 64e7d6ab9..000000000 Binary files a/docs/core_concept/imgs/grad_with_loss.png and /dev/null differ diff --git a/docs/core_concept/imgs/loss_with_net_and_rng.png b/docs/core_concept/imgs/loss_with_net_and_rng.png deleted file mode 100644 index 94e4b2af5..000000000 Binary files a/docs/core_concept/imgs/loss_with_net_and_rng.png and /dev/null differ diff --git a/docs/core_concept/imgs/train_with_grad_and_opt.png b/docs/core_concept/imgs/train_with_grad_and_opt.png deleted file mode 100644 index e5ff0ca30..000000000 Binary files a/docs/core_concept/imgs/train_with_grad_and_opt.png and /dev/null differ diff --git a/docs/index.rst b/docs/index.rst index 1cf3db2f3..583a30e08 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -195,8 +195,8 @@ Learn more core_concepts.rst tutorials.rst - advanced_tutorials.rst toolboxes.rst + advanced_tutorials.rst FAQ.rst api.rst diff --git a/docs/tutorial_FAQs/how_to_debug.ipynb b/docs/tutorial_FAQs/how_to_debug.ipynb index 8b51d0bc4..a6f124288 100644 --- a/docs/tutorial_FAQs/how_to_debug.ipynb +++ b/docs/tutorial_FAQs/how_to_debug.ipynb @@ -154,12 +154,11 @@ { "cell_type": "markdown", "source": [ - "``jax.disable_jit()`` works for all brainpy transformations, for example:\n", + "``jax.disable_jit()`` works for most brainpy transformations, including:\n", "\n", "- ``brainpy.math.jit()``\n", "- ``brainpy.math.grad()``\n", "- ``brainpy.math.vector_grad()``\n", - "- ``brainpy.math.for_loop()``\n", "- ``brainpy.math.while_loop()``\n", "- ``brainpy.math.cond()``\n", "- ``brainpy.math.ifelse()``" @@ -167,6 +166,28 @@ "metadata": { "collapsed": false } + }, + { + "cell_type": "markdown", + "source": [ + "## ``brainpy.DSRunner(..., jit=False)``\n", + "\n", + "If users are using ``brainpy.DSRunner``, you can initialize ``brainpy.DSRunner(..., jit=False)`` to disable JIT compilation when simulating a brain dynamics model.\n" + ], + "metadata": { + "collapsed": false + } + }, + { + "cell_type": "markdown", + "source": [ + "## ``brainpy.for_loop(..., jit=False)``\n", + "\n", + "Similarly, if users are using ``brainpy.for_loop``, you can put a ``jit=False`` argument into the ``for_loop`` transformation, then the JIT compilation will be removed." + ], + "metadata": { + "collapsed": false + } } ], "metadata": {