From 27e141d170549ff7c9d934cb5d098f1a811fd116 Mon Sep 17 00:00:00 2001 From: Sichao He <1310722434@qq.com> Date: Fri, 22 Nov 2024 17:29:44 +0800 Subject: [PATCH] Add lightweight approach checkpoint method of braintools (#40) --- docs/tutorials/checkpointing-en.ipynb | 285 +++++++++++++++++++++++-- docs/tutorials/checkpointing-zh.ipynb | 286 ++++++++++++++++++++++++-- 2 files changed, 543 insertions(+), 28 deletions(-) diff --git a/docs/tutorials/checkpointing-en.ipynb b/docs/tutorials/checkpointing-en.ipynb index 3af83bf..2cacf38 100644 --- a/docs/tutorials/checkpointing-en.ipynb +++ b/docs/tutorials/checkpointing-en.ipynb @@ -15,7 +15,7 @@ "id": "43f961f5", "metadata": {}, "source": [ - "In this tutorial, we will explore how to save and load checkpoints in `brainstate` by using the `orbax` library. This is particularly useful for saving the state of your model during training so that you can resume training from where you left off or use the trained model for inference later. The following example demonstrates how to use `orbax`'s checkpointing functionality with a simple MLP model." + "In this tutorial, we will explore how to save and load checkpoints in `brainstate` by using the `orbax` library and `braintools` library which provide a more lightweight approach. This is particularly useful for saving the state of your model during training so that you can resume training from where you left off or use the trained model for inference later. The following example demonstrates how to use `orbax` and `braintools`'s checkpointing functionality with a simple MLP model." ] }, { @@ -29,7 +29,11 @@ "\n", "You may also install directly from GitHub, using the following command. This can be used to obtain the most recent version of Orbax.\n", "\n", - "`pip install 'git+https://github.com/google/orbax/#subdirectory=checkpoint'`" + "`pip install 'git+https://github.com/google/orbax/#subdirectory=checkpoint'`\n", + "\n", + "You can install the `braintools` library by running the following command:\n", + "\n", + "`pip install braintools`" ] }, { @@ -42,7 +46,7 @@ }, { "cell_type": "code", - "execution_count": 9, + "execution_count": 14, "id": "b7741c32", "metadata": {}, "outputs": [], @@ -53,6 +57,7 @@ "import jax\n", "import jax.numpy as jnp\n", "import orbax.checkpoint as orbax\n", + "import braintools\n", "\n", "import brainstate as bst" ] @@ -68,7 +73,7 @@ }, { "cell_type": "code", - "execution_count": 8, + "execution_count": 15, "id": "7e020098", "metadata": {}, "outputs": [], @@ -97,7 +102,7 @@ }, { "cell_type": "code", - "execution_count": 11, + "execution_count": 16, "id": "39619169", "metadata": {}, "outputs": [ @@ -284,7 +289,7 @@ ")" ] }, - "execution_count": 11, + "execution_count": 16, "metadata": {}, "output_type": "execute_result" } @@ -302,12 +307,14 @@ "metadata": {}, "source": [ "## Save the Model State\n", + "\n", + "### Save the Model State Using `orbax`\n", "We save the model's parameters to a checkpoint file." ] }, { "cell_type": "code", - "execution_count": 15, + "execution_count": 17, "id": "14d1d552", "metadata": {}, "outputs": [], @@ -323,7 +330,42 @@ "id": "27209868", "metadata": {}, "source": [ - "Now, we've saved the model's parameters to the checkpoint files in `tmpdir/state`." + "Now, we've saved the model's parameters to the checkpoint files in `tmpdir/state` by using the `orbax` library." + ] + }, + { + "cell_type": "markdown", + "id": "fb36ffc3", + "metadata": {}, + "source": [ + "### Save the Model State Using `braintools`" + ] + }, + { + "cell_type": "code", + "execution_count": 18, + "id": "2b03606b", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Saving checkpoint into C:\\Users\\13107\\AppData\\Local\\Temp\\tmptjdpy0vf\\state.msgpack\n" + ] + } + ], + "source": [ + "checkpoint = bst.graph.states(model1).to_nest() # convert model to nest\n", + "braintools.file.msgpack_save(os.path.join(tmpdir, 'state.msgpack'), checkpoint) # save checkpoint" + ] + }, + { + "cell_type": "markdown", + "id": "76030ac1", + "metadata": {}, + "source": [ + "Now, we've saved the model's parameters to the checkpoint files in `tmpdir/state.msgpack` by using the `braintools` library." ] }, { @@ -332,12 +374,14 @@ "metadata": {}, "source": [ "## Load the Model State\n", + "\n", + "### Load the Model State Using `orbax`\n", "Let's load the model's parameters from the checkpoint files." ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 19, "id": "26ba3c3e", "metadata": {}, "outputs": [ @@ -345,7 +389,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "c:\\Users\\13107\\.conda\\envs\\brainpy-dev\\lib\\site-packages\\orbax\\checkpoint\\type_handlers.py:1386: UserWarning: Couldn't find sharding info under RestoreArgs. Populating sharding info from sharding file. Please note restoration time will be slightly increased due to reading from file instead of directly from RestoreArgs. Note also that this option is unsafe when restoring on a different topology than the checkpoint was saved with.\n", + "c:\\routhleck_app\\miniconda\\envs\\brainstate\\lib\\site-packages\\orbax\\checkpoint\\_src\\serialization\\type_handlers.py:1123: UserWarning: Couldn't find sharding info under RestoreArgs. Populating sharding info from sharding file. Please note restoration time will be slightly increased due to reading from file instead of directly from RestoreArgs. Note also that this option is unsafe when restoring on a different topology than the checkpoint was saved with.\n", " warnings.warn(\n" ] } @@ -364,6 +408,216 @@ "bst.graph.update_states(model2, state_tree)" ] }, + { + "cell_type": "markdown", + "id": "79929f4a", + "metadata": {}, + "source": [ + "### Load the Model State Using `braintools`\n", + "Let's load the model's parameters from the checkpoint files." + ] + }, + { + "cell_type": "code", + "execution_count": 20, + "id": "7a6d1de0", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Loading checkpoint from C:\\Users\\13107\\AppData\\Local\\Temp\\tmptjdpy0vf\\state.msgpack\n" + ] + }, + { + "data": { + "text/plain": [ + "{'dense1': {'weight': ParamState(\n", + " value={'bias': array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n", + " 0., 0., 0.], dtype=float32), 'weight': array([[ 0.74939334, 0.3148138 , 0.60089725, -0.7131149 , 0.6790908 ,\n", + " -0.44663328, 0.03113358, -0.5250644 , 0.1614144 , -0.39722365,\n", + " -0.23442519, 0.118144 , 0.7669531 , 0.06876656, 0.6045511 ,\n", + " 0.12086334, -0.88447595, -0.19188431, -0.85868365, 0.00500867],\n", + " [ 0.20412642, 0.07092498, 0.37392026, 0.34958398, -0.57214 ,\n", + " 0.71724516, -0.08160591, 0.50068825, -0.17175189, -0.08275215,\n", + " 0.6508336 , 0.28279537, 0.08821856, 0.83949256, 0.49844882,\n", + " -0.04159267, -0.47324428, 0.27084318, -0.58236146, -0.09787997],\n", + " [-0.04382031, -0.20300323, -0.04449642, 0.41578326, 0.5507486 ,\n", + " -0.15913244, -0.8612537 , 0.19072336, -0.16082875, -0.24696219,\n", + " -0.30372635, 0.6850187 , 0.32007053, 0.24253711, 0.28217098,\n", + " -0.8014343 , 0.48989874, -0.0160339 , 0.32790813, -0.49864978],\n", + " [-0.61840117, 0.21017133, 0.07593305, -0.02365256, -0.03401124,\n", + " -0.05115725, 0.6195931 , 0.15402867, 0.40200788, 0.34128165,\n", + " 0.00860781, -0.54993343, -0.5615623 , -0.09946032, -0.02702298,\n", + " 0.3336504 , -0.29341814, 0.3551176 , 0.20545702, -0.11665206],\n", + " [-0.16712527, -0.2531548 , 0.49188057, -0.1302325 , -0.12142995,\n", + " -0.03277557, 0.06477631, -0.30021554, -0.35658783, -0.5185722 ,\n", + " 0.15650164, -0.7464921 , -0.67454183, 0.09733332, -0.5153455 ,\n", + " 0.1480032 , -0.20877242, 0.16675173, 0.12827559, 0.5268865 ],\n", + " [-0.7994777 , -0.40662575, 0.28858158, -0.39780638, 0.6637344 ,\n", + " 0.09075797, -0.75130516, -0.26124355, 0.4175534 , -0.28502613,\n", + " -0.4241315 , 0.6746936 , 0.40870044, 0.94398546, -0.9198975 ,\n", + " -0.29775584, -0.09658122, -0.16053742, -0.05611025, 0.01059594],\n", + " [ 0.5480607 , -0.09164569, -0.7853424 , 0.74901533, -0.5906064 ,\n", + " -0.51409346, 0.10472732, -0.13107914, -0.45577446, -0.24654518,\n", + " 0.5399041 , -0.09071468, -0.5162382 , -0.01967659, -0.47176114,\n", + " -0.01017519, -0.5026951 , 0.05103482, 0.37542912, -0.25549397],\n", + " [-0.2706877 , 0.64187187, -0.505112 , -0.17481704, -0.88211423,\n", + " -0.8674219 , 0.5660908 , -0.20833156, 0.3285284 , 0.92883885,\n", + " -0.26592234, -0.47405127, 0.79681754, -0.5791843 , -0.27389136,\n", + " -0.3449671 , 0.509086 , 0.76971966, 0.10998839, -0.24425419],\n", + " [ 0.8046176 , -0.0295862 , 0.14252356, -0.1579972 , -0.20274054,\n", + " 0.01246137, -0.15756735, 0.32074738, 0.14097062, 0.03186554,\n", + " -0.1414449 , 0.4591949 , -0.21690284, -0.41089386, 0.26250118,\n", + " -0.0720875 , -0.05566718, -0.08271056, -0.37073353, 0.09257671],\n", + " [ 0.44894424, 0.22119072, -0.5117801 , -0.7407342 , -0.8777072 ,\n", + " 0.34723184, 0.0638053 , -0.10916334, 0.67356414, -0.21106955,\n", + " -0.24140975, 0.12431782, 0.2585294 , 0.06849731, -0.2997454 ,\n", + " -0.39390567, -0.25709096, -0.15120856, -0.10684931, 0.69015896]],\n", + " dtype=float32)}\n", + " )},\n", + " 'dense2': {'weight': ParamState(\n", + " value={'bias': array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n", + " 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32), 'weight': array([[ 0.62046814, 0.5301044 , -0.4739194 , -0.14099996, -0.14287984,\n", + " 0.1282555 , 0.3935479 , 0.21227883, 0.5402896 , -0.32984453,\n", + " 0.1054924 , -0.02015361, -0.24927817, 0.16467251, 0.5784846 ,\n", + " 0.2914683 , 0.35762057, 0.29866996, -0.19128309, 0.09088683,\n", + " -0.11386324, -0.22595015, 0.11267622, 0.5419977 , -0.37829107,\n", + " -0.09838869, -0.04575922, -0.7129366 , -0.32915255, 0.00509653],\n", + " [ 0.21818087, 0.08530099, 0.3571782 , 0.70128685, -0.04413987,\n", + " 0.5709911 , 0.12656331, 0.29721373, 0.47632915, -0.17275095,\n", + " -0.08733549, -0.22514656, -0.05714319, -0.27718347, -0.39045587,\n", + " 0.21975726, 0.18346666, -0.0382327 , 0.13839035, 0.1998283 ,\n", + " -0.09052311, 0.38183472, 0.4496051 , -0.23680712, -0.28785107,\n", + " 0.16122147, -0.33963904, 0.14983557, -0.43373275, -0.09495756],\n", + " [ 0.2568711 , 0.5197295 , -0.13442262, -0.6316247 , -0.6276094 ,\n", + " 0.396733 , 0.09978731, 0.37479848, 0.05811005, 0.38287428,\n", + " -0.23015432, 0.26524863, 0.40986276, 0.51085615, -0.16390967,\n", + " -0.08889349, 0.14242767, -0.04773026, 0.0186008 , 0.08168013,\n", + " 0.22218394, 0.45948145, 0.15798983, 0.11101982, -0.22625342,\n", + " -0.3179377 , 0.08289661, -0.35810882, -0.11701918, -0.07404168],\n", + " [ 0.3988777 , 0.09341867, -0.10675149, 0.24498817, 0.57484835,\n", + " 0.13964735, -0.09232395, 0.49800253, -0.11388287, -0.23314221,\n", + " -0.20017506, 0.17043568, -0.5916637 , -0.5033429 , -0.03982058,\n", + " -0.29196522, -0.06229761, -0.12120344, -0.04843295, 0.14077553,\n", + " -0.23975037, 0.25233614, -0.00446404, 0.6632397 , -0.32990777,\n", + " -0.42914438, -0.372548 , 0.30960974, 0.31027737, 0.3736987 ],\n", + " [-0.32519445, -0.0722824 , -0.06813759, 0.15726727, 0.52653533,\n", + " -0.39247712, -0.37830523, 0.20171025, -0.06937496, 0.24201019,\n", + " 0.1104718 , 0.62304336, 0.4803775 , -0.26503193, 0.5813743 ,\n", + " -0.22703817, 0.14889193, -0.09937828, 0.45811605, -0.53927666,\n", + " 0.38610622, 0.25877175, -0.57717675, -0.16893166, -0.17705517,\n", + " 0.2077132 , -0.24225888, -0.11191322, -0.00921882, -0.10405794],\n", + " [ 0.41278893, -0.27192885, 0.28467888, -0.21523082, 0.37667713,\n", + " 0.07426698, 0.22414407, -0.1354481 , -0.23419291, 0.2381074 ,\n", + " -0.24765436, 0.08778596, -0.00406975, -0.615931 , -0.09067997,\n", + " 0.26324016, -0.03728105, 0.29038942, 0.678011 , -0.6540893 ,\n", + " -0.5934551 , -0.16575795, 0.14227462, -0.0928836 , 0.24194399,\n", + " -0.04459891, 0.15232474, -0.21208623, -0.21339062, 0.07757895],\n", + " [-0.6379539 , 0.31518504, -0.11890189, -0.19096668, 0.21524261,\n", + " -0.06361473, 0.56184316, 0.028249 , -0.14510861, 0.08830918,\n", + " 0.08343762, -0.25384745, -0.33789673, 0.03700592, -0.19126455,\n", + " -0.01024354, -0.37079507, 0.24292567, 0.19478266, 0.5580041 ,\n", + " -0.35604435, 0.3915089 , -0.21796615, 0.0528199 , -0.13147084,\n", + " -0.05164728, -0.0625616 , 0.36192182, -0.05759151, 0.4186158 ],\n", + " [-0.04047865, 0.02108607, 0.41284686, 0.29146758, 0.20885086,\n", + " 0.20158692, -0.17301778, 0.27862224, 0.27474535, -0.19628745,\n", + " 0.15615414, 0.20871529, -0.314695 , -0.24115679, 0.33787283,\n", + " -0.14589988, -0.10813709, -0.039655 , -0.03082952, -0.66367936,\n", + " -0.2642637 , 0.2510051 , -0.08893799, 0.21589737, 0.51835227,\n", + " -0.44741842, -0.33786973, 0.6091706 , -0.3753065 , -0.37535354],\n", + " [ 0.11531412, 0.6267082 , -0.15149857, -0.3794238 , 0.55059415,\n", + " 0.23017633, -0.32434496, 0.2958217 , 0.41106105, 0.4731116 ,\n", + " -0.50055134, 0.01790522, -0.54518443, 0.04447998, -0.13089894,\n", + " -0.15774457, 0.09551436, -0.08697572, -0.05562068, -0.06885753,\n", + " 0.20314606, 0.14044988, -0.19203717, -0.4179157 , 0.18612123,\n", + " -0.14104603, -0.35670066, -0.24597271, 0.10614085, -0.12170368],\n", + " [ 0.23700227, 0.30524203, -0.3694181 , 0.33033338, 0.02095676,\n", + " -0.05125551, 0.11001365, -0.20992021, -0.05562193, -0.26372904,\n", + " -0.2967057 , -0.14012977, -0.14321879, -0.17379181, 0.5104145 ,\n", + " 0.11991877, -0.1430745 , -0.04331772, -0.41226274, 0.00449552,\n", + " -0.08277246, -0.12151891, -0.45340443, 0.12951623, -0.27139285,\n", + " 0.4472014 , 0.19157353, -0.4412653 , -0.04408614, 0.41542286],\n", + " [ 0.04913985, -0.04957955, -0.40214545, -0.24126607, -0.11509801,\n", + " -0.51304626, -0.3825655 , 0.34506062, -0.0222565 , -0.27472144,\n", + " -0.5477002 , -0.03630246, 0.17396483, 0.6892827 , 0.02867843,\n", + " 0.36273733, -0.34478036, 0.2839792 , 0.15002191, -0.20483544,\n", + " 0.15306501, -0.06504299, -0.00701311, 0.0804052 , 0.44663915,\n", + " 0.11938784, -0.05011488, 0.06942522, -0.1151372 , 0.2728172 ],\n", + " [-0.30464825, 0.11323573, 0.02953907, -0.7024937 , -0.04522578,\n", + " 0.10622236, -0.1298965 , 0.0872021 , -0.36016473, -0.11690426,\n", + " -0.07054564, -0.32576308, -0.30710763, -0.6661573 , 0.13130474,\n", + " 0.00769307, 0.00603968, -0.5331483 , -0.00946458, -0.08804175,\n", + " 0.01258891, 0.19920264, -0.52920264, 0.11547033, 0.0503376 ,\n", + " 0.2710771 , 0.20577058, -0.16118994, 0.03479335, 0.30332327],\n", + " [-0.11540684, -0.21528308, -0.09639532, -0.38324118, 0.08790598,\n", + " -0.05113763, -0.22907412, 0.08176684, -0.13504112, -0.14580515,\n", + " -0.10574839, -0.13816664, 0.25279123, -0.35016036, -0.02811426,\n", + " 0.1878024 , 0.33833987, -0.44787505, 0.05859555, -0.12482259,\n", + " 0.4109398 , -0.3567587 , 0.4436607 , -0.13256377, 0.42250675,\n", + " 0.33017033, 0.28086263, 0.33791474, 0.24015151, -0.23016477],\n", + " [ 0.46682912, -0.63216 , 0.43159592, 0.21971288, -0.07587896,\n", + " -0.25639635, -0.42970398, -0.4962936 , -0.21198583, 0.18351796,\n", + " 0.01911162, -0.3004833 , -0.41785267, -0.04077749, -0.20676233,\n", + " -0.11401828, 0.12992048, 0.03491049, 0.05013497, 0.57222587,\n", + " -0.12001502, -0.17038153, -0.31871405, -0.32121637, 0.66278815,\n", + " 0.61774564, -0.01240813, -0.06011448, 0.29245874, -0.3879291 ],\n", + " [ 0.02741514, 0.31249774, -0.15944321, 0.14222006, 0.611036 ,\n", + " 0.02716783, 0.48367155, -0.59191144, -0.260246 , 0.29856846,\n", + " 0.36217022, 0.26721174, 0.1436277 , 0.2510483 , 0.63455343,\n", + " 0.22804502, 0.21089312, -0.03622444, 0.24770333, 0.12762095,\n", + " -0.11348359, 0.71003526, -0.6399693 , 0.2956937 , -0.40721762,\n", + " 0.07830685, -0.12750737, 0.09320084, -0.37348104, 0.6469367 ],\n", + " [-0.21946031, 0.58491176, 0.6910229 , -0.38729444, -0.22691855,\n", + " 0.09827446, -0.27745098, 0.3286477 , -0.28397417, 0.3331472 ,\n", + " -0.10511833, 0.04856022, 0.6826674 , -0.19410591, -0.03848339,\n", + " 0.2877471 , 0.42053938, -0.3121656 , 0.1115057 , 0.3940428 ,\n", + " 0.22287792, -0.11617415, -0.15520288, -0.17891021, 0.08283449,\n", + " -0.45727572, -0.08755263, -0.30042952, 0.04397725, -0.32858402],\n", + " [-0.04652168, 0.22256051, 0.34796244, -0.57714033, -0.19478762,\n", + " -0.04000793, -0.22230573, -0.1784827 , 0.18552966, 0.3517072 ,\n", + " -0.43350866, 0.3370349 , 0.34543782, -0.25484002, -0.06113737,\n", + " -0.29600585, 0.55229264, 0.26264954, -0.12024187, 0.06554315,\n", + " 0.33039162, -0.4056347 , -0.22326599, -0.20423931, -0.20365807,\n", + " 0.5614395 , -0.33278635, 0.3678192 , -0.38601917, -0.12349749],\n", + " [ 0.21260151, -0.6383393 , -0.04182729, 0.21110533, -0.16549559,\n", + " 0.20241106, 0.42155504, 0.2782736 , -0.5695076 , 0.3197464 ,\n", + " 0.3593777 , 0.15281492, -0.16649725, -0.32258078, -0.19450592,\n", + " -0.5648749 , 0.14112377, -0.08617025, 0.2822599 , 0.65894157,\n", + " 0.06424519, -0.02703291, 0.41351956, -0.06962998, -0.03156902,\n", + " -0.3027034 , -0.15010884, 0.3097132 , -0.01670518, 0.13812247],\n", + " [-0.35231128, -0.06400244, -0.5534636 , 0.08153537, -0.1431605 ,\n", + " 0.19649687, -0.57627857, 0.14731233, -0.5345133 , 0.14830953,\n", + " 0.11090186, -0.5130216 , 0.07951056, 0.042261 , 0.0088584 ,\n", + " 0.0693031 , -0.25705618, 0.07637526, -0.2910843 , 0.26884285,\n", + " -0.3668523 , -0.51732624, 0.32633176, 0.4078384 , 0.07319385,\n", + " 0.24243955, -0.39059573, -0.14434972, -0.20902094, 0.03081408],\n", + " [-0.29074088, -0.340606 , 0.24403909, 0.28382063, 0.57466537,\n", + " 0.24103518, -0.53504395, -0.12040613, -0.21954668, -0.11855581,\n", + " 0.20805535, -0.6497588 , 0.03112273, -0.06355662, 0.22711465,\n", + " -0.00476316, -0.4368407 , -0.26775414, 0.02075309, -0.0473614 ,\n", + " -0.12880138, 0.15983032, 0.18893135, -0.06872427, -0.14535248,\n", + " 0.27104148, -0.31298438, 0.14454837, -0.1837953 , 0.4652801 ]],\n", + " dtype=float32)}\n", + " )}}" + ] + }, + "execution_count": 20, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# 创建一个有着相同结构的模型\n", + "bst.random.seed(0)\n", + "model3 = bst.augment.eval_shape(lambda: MLP(10, 20, 30))\n", + "checkpoint = bst.graph.states(model3).to_nest()\n", + "\n", + "# 从msgpack文件读取模型参数\n", + "braintools.file.msgpack_load(os.path.join(tmpdir, 'state.msgpack'), checkpoint)" + ] + }, { "cell_type": "markdown", "id": "29dc37c9", @@ -375,7 +629,7 @@ }, { "cell_type": "code", - "execution_count": 18, + "execution_count": 21, "id": "dfe032ab", "metadata": {}, "outputs": [ @@ -383,6 +637,7 @@ "name": "stdout", "output_type": "stream", "text": [ + "True\n", "True\n" ] } @@ -390,13 +645,15 @@ "source": [ "y1 = model1(jnp.ones((1, 10)))\n", "y2 = model2(jnp.ones((1, 10)))\n", - "print(jnp.allclose(y1, y2)) # True" + "y3 = model3(jnp.ones((1, 10)))\n", + "print(jnp.allclose(y1, y2)) # True\n", + "print(jnp.allclose(y1, y3)) # True" ] } ], "metadata": { "kernelspec": { - "display_name": "brainpy-dev", + "display_name": "brainstate", "language": "python", "name": "python3" }, @@ -410,7 +667,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.13" + "version": "3.10.15" } }, "nbformat": 4, diff --git a/docs/tutorials/checkpointing-zh.ipynb b/docs/tutorials/checkpointing-zh.ipynb index 95b7fe0..c050ccc 100644 --- a/docs/tutorials/checkpointing-zh.ipynb +++ b/docs/tutorials/checkpointing-zh.ipynb @@ -15,7 +15,7 @@ "id": "d775430d", "metadata": {}, "source": [ - "在本教程中，我们将探讨如何使用`orbax`库在`brainstate`中保存和加载检查点。这对于在训练过程中保存模型状态非常有用，这样您可以从中断的地方继续训练或稍后使用已训练的模型进行推理。以下示例演示了如何将`orbax`的检查点功能与一个简单的多层感知机(MLP)模型结合使用。" + "在本教程中，我们将探讨如何使用`orbax`库以及`braintools`的轻量级方法在`brainstate`中保存和加载检查点。这对于在训练过程中保存模型状态非常有用，这样您可以从中断的地方继续训练或稍后使用已训练的模型进行推理。以下示例演示了如何将`orbax`和`braintools`的检查点功能与一个简单的多层感知机(MLP)模型结合使用。" ] }, { @@ -29,7 +29,11 @@ "\n", "您也可以直接从 GitHub 安装，使用以下命令。这可以用来获取 Orbax 的最新版本。\n", "\n", - "`pip install 'git+https://github.com/google/orbax/#subdirectory=checkpoint'`" + "`pip install 'git+https://github.com/google/orbax/#subdirectory=checkpoint'`\n", + "\n", + "其次，您可以通过运行以下命令安装`braintools`库：\n", + "\n", + "`pip install braintools`" ] }, { @@ -42,7 +46,7 @@ }, { "cell_type": "code", - "execution_count": 1, + "execution_count": 18, "id": "8142091f", "metadata": {}, "outputs": [], @@ -53,6 +57,7 @@ "import jax\n", "import jax.numpy as jnp\n", "import orbax.checkpoint as orbax\n", + "import braintools\n", "\n", "import brainstate as bst" ] @@ -68,7 +73,7 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": 19, "id": "d090d2e3", "metadata": {}, "outputs": [], @@ -97,7 +102,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 20, "id": "67ca04d3", "metadata": {}, "outputs": [ @@ -284,8 +289,9 @@ ")" ] }, + "execution_count": 20, "metadata": {}, - "output_type": "display_data" + "output_type": "execute_result" } ], "source": [ @@ -301,12 +307,14 @@ "metadata": {}, "source": [ "## 保存模型参数\n", + "\n", + "### 使用`orbax`保存检查点\n", "我们将模型参数保存到检查点文件中。" ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 21, "id": "c529099d", "metadata": {}, "outputs": [], @@ -322,7 +330,42 @@ "id": "92448b34", "metadata": {}, "source": [ - "现在，我们已经将模型的参数保存到`tmpdir/state`的检查点文件中。" + "现在，我们已经将模型的参数通过`orbax`保存到`tmpdir/state`的检查点文件中。" + ] + }, + { + "cell_type": "markdown", + "id": "a8b032ba", + "metadata": {}, + "source": [ + "### 使用`braintools`保存检查点" + ] + }, + { + "cell_type": "code", + "execution_count": 22, + "id": "ed7f1c6e", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Saving checkpoint into C:\\Users\\13107\\AppData\\Local\\Temp\\tmp483fc4t1\\state.msgpack\n" + ] + } + ], + "source": [ + "checkpoint = bst.graph.states(model1).to_nest() # 将模型的状态转换为nest结构\n", + "braintools.file.msgpack_save(os.path.join(tmpdir, 'state.msgpack'), checkpoint) # 保存模型的参数" + ] + }, + { + "cell_type": "markdown", + "id": "18976368", + "metadata": {}, + "source": [ + "现在，我们已经将模型的参数通过`braintools`保存到`tmpdir/state.msgpack`的检查点文件中。" ] }, { @@ -331,12 +374,14 @@ "metadata": {}, "source": [ "## 加载模型参数\n", + "\n", + "### 使用`orbax`加载检查点\n", "我们将从检查点文件中加载模型的参数。" ] }, { "cell_type": "code", - "execution_count": null, + "execution_count": 23, "id": "db238309", "metadata": {}, "outputs": [ @@ -344,7 +389,7 @@ "name": "stderr", "output_type": "stream", "text": [ - "c:\\Users\\13107\\.conda\\envs\\brainpy-dev\\lib\\site-packages\\orbax\\checkpoint\\type_handlers.py:1386: UserWarning: Couldn't find sharding info under RestoreArgs. Populating sharding info from sharding file. Please note restoration time will be slightly increased due to reading from file instead of directly from RestoreArgs. Note also that this option is unsafe when restoring on a different topology than the checkpoint was saved with.\n", + "c:\\routhleck_app\\miniconda\\envs\\brainstate\\lib\\site-packages\\orbax\\checkpoint\\_src\\serialization\\type_handlers.py:1123: UserWarning: Couldn't find sharding info under RestoreArgs. Populating sharding info from sharding file. Please note restoration time will be slightly increased due to reading from file instead of directly from RestoreArgs. Note also that this option is unsafe when restoring on a different topology than the checkpoint was saved with.\n", " warnings.warn(\n" ] } @@ -363,6 +408,216 @@ "bst.graph.update_states(model2, state_tree)" ] }, + { + "cell_type": "markdown", + "id": "aa75ceb7", + "metadata": {}, + "source": [ + "### 使用`braintools`加载检查点\n", + "我们将从检查点文件中加载模型的参数。" + ] + }, + { + "cell_type": "code", + "execution_count": 24, + "id": "b0fbd6f6", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Loading checkpoint from C:\\Users\\13107\\AppData\\Local\\Temp\\tmp483fc4t1\\state.msgpack\n" + ] + }, + { + "data": { + "text/plain": [ + "{'dense1': {'weight': ParamState(\n", + " value={'bias': array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n", + " 0., 0., 0.], dtype=float32), 'weight': array([[ 0.74939334, 0.3148138 , 0.60089725, -0.7131149 , 0.6790908 ,\n", + " -0.44663328, 0.03113358, -0.5250644 , 0.1614144 , -0.39722365,\n", + " -0.23442519, 0.118144 , 0.7669531 , 0.06876656, 0.6045511 ,\n", + " 0.12086334, -0.88447595, -0.19188431, -0.85868365, 0.00500867],\n", + " [ 0.20412642, 0.07092498, 0.37392026, 0.34958398, -0.57214 ,\n", + " 0.71724516, -0.08160591, 0.50068825, -0.17175189, -0.08275215,\n", + " 0.6508336 , 0.28279537, 0.08821856, 0.83949256, 0.49844882,\n", + " -0.04159267, -0.47324428, 0.27084318, -0.58236146, -0.09787997],\n", + " [-0.04382031, -0.20300323, -0.04449642, 0.41578326, 0.5507486 ,\n", + " -0.15913244, -0.8612537 , 0.19072336, -0.16082875, -0.24696219,\n", + " -0.30372635, 0.6850187 , 0.32007053, 0.24253711, 0.28217098,\n", + " -0.8014343 , 0.48989874, -0.0160339 , 0.32790813, -0.49864978],\n", + " [-0.61840117, 0.21017133, 0.07593305, -0.02365256, -0.03401124,\n", + " -0.05115725, 0.6195931 , 0.15402867, 0.40200788, 0.34128165,\n", + " 0.00860781, -0.54993343, -0.5615623 , -0.09946032, -0.02702298,\n", + " 0.3336504 , -0.29341814, 0.3551176 , 0.20545702, -0.11665206],\n", + " [-0.16712527, -0.2531548 , 0.49188057, -0.1302325 , -0.12142995,\n", + " -0.03277557, 0.06477631, -0.30021554, -0.35658783, -0.5185722 ,\n", + " 0.15650164, -0.7464921 , -0.67454183, 0.09733332, -0.5153455 ,\n", + " 0.1480032 , -0.20877242, 0.16675173, 0.12827559, 0.5268865 ],\n", + " [-0.7994777 , -0.40662575, 0.28858158, -0.39780638, 0.6637344 ,\n", + " 0.09075797, -0.75130516, -0.26124355, 0.4175534 , -0.28502613,\n", + " -0.4241315 , 0.6746936 , 0.40870044, 0.94398546, -0.9198975 ,\n", + " -0.29775584, -0.09658122, -0.16053742, -0.05611025, 0.01059594],\n", + " [ 0.5480607 , -0.09164569, -0.7853424 , 0.74901533, -0.5906064 ,\n", + " -0.51409346, 0.10472732, -0.13107914, -0.45577446, -0.24654518,\n", + " 0.5399041 , -0.09071468, -0.5162382 , -0.01967659, -0.47176114,\n", + " -0.01017519, -0.5026951 , 0.05103482, 0.37542912, -0.25549397],\n", + " [-0.2706877 , 0.64187187, -0.505112 , -0.17481704, -0.88211423,\n", + " -0.8674219 , 0.5660908 , -0.20833156, 0.3285284 , 0.92883885,\n", + " -0.26592234, -0.47405127, 0.79681754, -0.5791843 , -0.27389136,\n", + " -0.3449671 , 0.509086 , 0.76971966, 0.10998839, -0.24425419],\n", + " [ 0.8046176 , -0.0295862 , 0.14252356, -0.1579972 , -0.20274054,\n", + " 0.01246137, -0.15756735, 0.32074738, 0.14097062, 0.03186554,\n", + " -0.1414449 , 0.4591949 , -0.21690284, -0.41089386, 0.26250118,\n", + " -0.0720875 , -0.05566718, -0.08271056, -0.37073353, 0.09257671],\n", + " [ 0.44894424, 0.22119072, -0.5117801 , -0.7407342 , -0.8777072 ,\n", + " 0.34723184, 0.0638053 , -0.10916334, 0.67356414, -0.21106955,\n", + " -0.24140975, 0.12431782, 0.2585294 , 0.06849731, -0.2997454 ,\n", + " -0.39390567, -0.25709096, -0.15120856, -0.10684931, 0.69015896]],\n", + " dtype=float32)}\n", + " )},\n", + " 'dense2': {'weight': ParamState(\n", + " value={'bias': array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n", + " 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float32), 'weight': array([[ 0.62046814, 0.5301044 , -0.4739194 , -0.14099996, -0.14287984,\n", + " 0.1282555 , 0.3935479 , 0.21227883, 0.5402896 , -0.32984453,\n", + " 0.1054924 , -0.02015361, -0.24927817, 0.16467251, 0.5784846 ,\n", + " 0.2914683 , 0.35762057, 0.29866996, -0.19128309, 0.09088683,\n", + " -0.11386324, -0.22595015, 0.11267622, 0.5419977 , -0.37829107,\n", + " -0.09838869, -0.04575922, -0.7129366 , -0.32915255, 0.00509653],\n", + " [ 0.21818087, 0.08530099, 0.3571782 , 0.70128685, -0.04413987,\n", + " 0.5709911 , 0.12656331, 0.29721373, 0.47632915, -0.17275095,\n", + " -0.08733549, -0.22514656, -0.05714319, -0.27718347, -0.39045587,\n", + " 0.21975726, 0.18346666, -0.0382327 , 0.13839035, 0.1998283 ,\n", + " -0.09052311, 0.38183472, 0.4496051 , -0.23680712, -0.28785107,\n", + " 0.16122147, -0.33963904, 0.14983557, -0.43373275, -0.09495756],\n", + " [ 0.2568711 , 0.5197295 , -0.13442262, -0.6316247 , -0.6276094 ,\n", + " 0.396733 , 0.09978731, 0.37479848, 0.05811005, 0.38287428,\n", + " -0.23015432, 0.26524863, 0.40986276, 0.51085615, -0.16390967,\n", + " -0.08889349, 0.14242767, -0.04773026, 0.0186008 , 0.08168013,\n", + " 0.22218394, 0.45948145, 0.15798983, 0.11101982, -0.22625342,\n", + " -0.3179377 , 0.08289661, -0.35810882, -0.11701918, -0.07404168],\n", + " [ 0.3988777 , 0.09341867, -0.10675149, 0.24498817, 0.57484835,\n", + " 0.13964735, -0.09232395, 0.49800253, -0.11388287, -0.23314221,\n", + " -0.20017506, 0.17043568, -0.5916637 , -0.5033429 , -0.03982058,\n", + " -0.29196522, -0.06229761, -0.12120344, -0.04843295, 0.14077553,\n", + " -0.23975037, 0.25233614, -0.00446404, 0.6632397 , -0.32990777,\n", + " -0.42914438, -0.372548 , 0.30960974, 0.31027737, 0.3736987 ],\n", + " [-0.32519445, -0.0722824 , -0.06813759, 0.15726727, 0.52653533,\n", + " -0.39247712, -0.37830523, 0.20171025, -0.06937496, 0.24201019,\n", + " 0.1104718 , 0.62304336, 0.4803775 , -0.26503193, 0.5813743 ,\n", + " -0.22703817, 0.14889193, -0.09937828, 0.45811605, -0.53927666,\n", + " 0.38610622, 0.25877175, -0.57717675, -0.16893166, -0.17705517,\n", + " 0.2077132 , -0.24225888, -0.11191322, -0.00921882, -0.10405794],\n", + " [ 0.41278893, -0.27192885, 0.28467888, -0.21523082, 0.37667713,\n", + " 0.07426698, 0.22414407, -0.1354481 , -0.23419291, 0.2381074 ,\n", + " -0.24765436, 0.08778596, -0.00406975, -0.615931 , -0.09067997,\n", + " 0.26324016, -0.03728105, 0.29038942, 0.678011 , -0.6540893 ,\n", + " -0.5934551 , -0.16575795, 0.14227462, -0.0928836 , 0.24194399,\n", + " -0.04459891, 0.15232474, -0.21208623, -0.21339062, 0.07757895],\n", + " [-0.6379539 , 0.31518504, -0.11890189, -0.19096668, 0.21524261,\n", + " -0.06361473, 0.56184316, 0.028249 , -0.14510861, 0.08830918,\n", + " 0.08343762, -0.25384745, -0.33789673, 0.03700592, -0.19126455,\n", + " -0.01024354, -0.37079507, 0.24292567, 0.19478266, 0.5580041 ,\n", + " -0.35604435, 0.3915089 , -0.21796615, 0.0528199 , -0.13147084,\n", + " -0.05164728, -0.0625616 , 0.36192182, -0.05759151, 0.4186158 ],\n", + " [-0.04047865, 0.02108607, 0.41284686, 0.29146758, 0.20885086,\n", + " 0.20158692, -0.17301778, 0.27862224, 0.27474535, -0.19628745,\n", + " 0.15615414, 0.20871529, -0.314695 , -0.24115679, 0.33787283,\n", + " -0.14589988, -0.10813709, -0.039655 , -0.03082952, -0.66367936,\n", + " -0.2642637 , 0.2510051 , -0.08893799, 0.21589737, 0.51835227,\n", + " -0.44741842, -0.33786973, 0.6091706 , -0.3753065 , -0.37535354],\n", + " [ 0.11531412, 0.6267082 , -0.15149857, -0.3794238 , 0.55059415,\n", + " 0.23017633, -0.32434496, 0.2958217 , 0.41106105, 0.4731116 ,\n", + " -0.50055134, 0.01790522, -0.54518443, 0.04447998, -0.13089894,\n", + " -0.15774457, 0.09551436, -0.08697572, -0.05562068, -0.06885753,\n", + " 0.20314606, 0.14044988, -0.19203717, -0.4179157 , 0.18612123,\n", + " -0.14104603, -0.35670066, -0.24597271, 0.10614085, -0.12170368],\n", + " [ 0.23700227, 0.30524203, -0.3694181 , 0.33033338, 0.02095676,\n", + " -0.05125551, 0.11001365, -0.20992021, -0.05562193, -0.26372904,\n", + " -0.2967057 , -0.14012977, -0.14321879, -0.17379181, 0.5104145 ,\n", + " 0.11991877, -0.1430745 , -0.04331772, -0.41226274, 0.00449552,\n", + " -0.08277246, -0.12151891, -0.45340443, 0.12951623, -0.27139285,\n", + " 0.4472014 , 0.19157353, -0.4412653 , -0.04408614, 0.41542286],\n", + " [ 0.04913985, -0.04957955, -0.40214545, -0.24126607, -0.11509801,\n", + " -0.51304626, -0.3825655 , 0.34506062, -0.0222565 , -0.27472144,\n", + " -0.5477002 , -0.03630246, 0.17396483, 0.6892827 , 0.02867843,\n", + " 0.36273733, -0.34478036, 0.2839792 , 0.15002191, -0.20483544,\n", + " 0.15306501, -0.06504299, -0.00701311, 0.0804052 , 0.44663915,\n", + " 0.11938784, -0.05011488, 0.06942522, -0.1151372 , 0.2728172 ],\n", + " [-0.30464825, 0.11323573, 0.02953907, -0.7024937 , -0.04522578,\n", + " 0.10622236, -0.1298965 , 0.0872021 , -0.36016473, -0.11690426,\n", + " -0.07054564, -0.32576308, -0.30710763, -0.6661573 , 0.13130474,\n", + " 0.00769307, 0.00603968, -0.5331483 , -0.00946458, -0.08804175,\n", + " 0.01258891, 0.19920264, -0.52920264, 0.11547033, 0.0503376 ,\n", + " 0.2710771 , 0.20577058, -0.16118994, 0.03479335, 0.30332327],\n", + " [-0.11540684, -0.21528308, -0.09639532, -0.38324118, 0.08790598,\n", + " -0.05113763, -0.22907412, 0.08176684, -0.13504112, -0.14580515,\n", + " -0.10574839, -0.13816664, 0.25279123, -0.35016036, -0.02811426,\n", + " 0.1878024 , 0.33833987, -0.44787505, 0.05859555, -0.12482259,\n", + " 0.4109398 , -0.3567587 , 0.4436607 , -0.13256377, 0.42250675,\n", + " 0.33017033, 0.28086263, 0.33791474, 0.24015151, -0.23016477],\n", + " [ 0.46682912, -0.63216 , 0.43159592, 0.21971288, -0.07587896,\n", + " -0.25639635, -0.42970398, -0.4962936 , -0.21198583, 0.18351796,\n", + " 0.01911162, -0.3004833 , -0.41785267, -0.04077749, -0.20676233,\n", + " -0.11401828, 0.12992048, 0.03491049, 0.05013497, 0.57222587,\n", + " -0.12001502, -0.17038153, -0.31871405, -0.32121637, 0.66278815,\n", + " 0.61774564, -0.01240813, -0.06011448, 0.29245874, -0.3879291 ],\n", + " [ 0.02741514, 0.31249774, -0.15944321, 0.14222006, 0.611036 ,\n", + " 0.02716783, 0.48367155, -0.59191144, -0.260246 , 0.29856846,\n", + " 0.36217022, 0.26721174, 0.1436277 , 0.2510483 , 0.63455343,\n", + " 0.22804502, 0.21089312, -0.03622444, 0.24770333, 0.12762095,\n", + " -0.11348359, 0.71003526, -0.6399693 , 0.2956937 , -0.40721762,\n", + " 0.07830685, -0.12750737, 0.09320084, -0.37348104, 0.6469367 ],\n", + " [-0.21946031, 0.58491176, 0.6910229 , -0.38729444, -0.22691855,\n", + " 0.09827446, -0.27745098, 0.3286477 , -0.28397417, 0.3331472 ,\n", + " -0.10511833, 0.04856022, 0.6826674 , -0.19410591, -0.03848339,\n", + " 0.2877471 , 0.42053938, -0.3121656 , 0.1115057 , 0.3940428 ,\n", + " 0.22287792, -0.11617415, -0.15520288, -0.17891021, 0.08283449,\n", + " -0.45727572, -0.08755263, -0.30042952, 0.04397725, -0.32858402],\n", + " [-0.04652168, 0.22256051, 0.34796244, -0.57714033, -0.19478762,\n", + " -0.04000793, -0.22230573, -0.1784827 , 0.18552966, 0.3517072 ,\n", + " -0.43350866, 0.3370349 , 0.34543782, -0.25484002, -0.06113737,\n", + " -0.29600585, 0.55229264, 0.26264954, -0.12024187, 0.06554315,\n", + " 0.33039162, -0.4056347 , -0.22326599, -0.20423931, -0.20365807,\n", + " 0.5614395 , -0.33278635, 0.3678192 , -0.38601917, -0.12349749],\n", + " [ 0.21260151, -0.6383393 , -0.04182729, 0.21110533, -0.16549559,\n", + " 0.20241106, 0.42155504, 0.2782736 , -0.5695076 , 0.3197464 ,\n", + " 0.3593777 , 0.15281492, -0.16649725, -0.32258078, -0.19450592,\n", + " -0.5648749 , 0.14112377, -0.08617025, 0.2822599 , 0.65894157,\n", + " 0.06424519, -0.02703291, 0.41351956, -0.06962998, -0.03156902,\n", + " -0.3027034 , -0.15010884, 0.3097132 , -0.01670518, 0.13812247],\n", + " [-0.35231128, -0.06400244, -0.5534636 , 0.08153537, -0.1431605 ,\n", + " 0.19649687, -0.57627857, 0.14731233, -0.5345133 , 0.14830953,\n", + " 0.11090186, -0.5130216 , 0.07951056, 0.042261 , 0.0088584 ,\n", + " 0.0693031 , -0.25705618, 0.07637526, -0.2910843 , 0.26884285,\n", + " -0.3668523 , -0.51732624, 0.32633176, 0.4078384 , 0.07319385,\n", + " 0.24243955, -0.39059573, -0.14434972, -0.20902094, 0.03081408],\n", + " [-0.29074088, -0.340606 , 0.24403909, 0.28382063, 0.57466537,\n", + " 0.24103518, -0.53504395, -0.12040613, -0.21954668, -0.11855581,\n", + " 0.20805535, -0.6497588 , 0.03112273, -0.06355662, 0.22711465,\n", + " -0.00476316, -0.4368407 , -0.26775414, 0.02075309, -0.0473614 ,\n", + " -0.12880138, 0.15983032, 0.18893135, -0.06872427, -0.14535248,\n", + " 0.27104148, -0.31298438, 0.14454837, -0.1837953 , 0.4652801 ]],\n", + " dtype=float32)}\n", + " )}}" + ] + }, + "execution_count": 24, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# 创建一个有着相同结构的模型\n", + "bst.random.seed(0)\n", + "model3 = bst.augment.eval_shape(lambda: MLP(10, 20, 30))\n", + "checkpoint = bst.graph.states(model3).to_nest()\n", + "\n", + "# 从msgpack文件读取模型参数\n", + "braintools.file.msgpack_load(os.path.join(tmpdir, 'state.msgpack'), checkpoint)" + ] + }, { "cell_type": "markdown", "id": "11994171", @@ -374,7 +629,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 25, "id": "64810e26", "metadata": {}, "outputs": [ @@ -382,6 +637,7 @@ "name": "stdout", "output_type": "stream", "text": [ + "True\n", "True\n" ] } @@ -389,13 +645,15 @@ "source": [ "y1 = model1(jnp.ones((1, 10)))\n", "y2 = model2(jnp.ones((1, 10)))\n", - "print(jnp.allclose(y1, y2)) # True" + "y3 = model3(jnp.ones((1, 10)))\n", + "print(jnp.allclose(y1, y2)) # True\n", + "print(jnp.allclose(y1, y3)) # True" ] } ], "metadata": { "kernelspec": { - "display_name": "brainpy-dev", + "display_name": "brainstate", "language": "python", "name": "python3" }, @@ -409,7 +667,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.13" + "version": "3.10.15" } }, "nbformat": 4,