prepare_training_with_decay_simple.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.core.interactiveshell import InteractiveShell\n",
    "\n",
    "InteractiveShell.ast_node_interactivity = \"all\"\n",
    "%reload_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from flax import nnx\n",
    "\n",
    "\n",
    "class Model(nnx.Module):\n",
    "    def __init__(self, rngs):\n",
    "        self.linear1 = nnx.Linear(2, 3, rngs=rngs)\n",
    "        self.linear2 = nnx.Linear(3, 4, rngs=rngs)\n",
    "\n",
    "    def __call__(self, x):\n",
    "        return self.linear2(self.linear1(x))\n",
    "\n",
    "\n",
    "x = jax.random.normal(jax.random.key(0), (1, 2))\n",
    "y = jnp.ones((1, 4))\n",
    "\n",
    "model = Model(nnx.Rngs(0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "State({\n",
       "  'linear1': {\n",
       "    'bias': False,\n",
       "    'kernel': True\n",
       "  },\n",
       "  'linear2': {\n",
       "    'bias': False,\n",
       "    'kernel': True\n",
       "  }\n",
       "})"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from flax import nnx\n",
    "\n",
    "\n",
    "def model_state_decay_mask(model):\n",
    "    flat_state = nnx.state(model).flat_state()\n",
    "    flat_mask = flat_state.copy()\n",
    "    for key in flat_state.keys():\n",
    "        flat_mask[key] = key[-1] not in (\"bias\")\n",
    "    # return frozen_dict.freeze(nnx.State.from_flat_path(flat_mask))\n",
    "    return nnx.State.from_flat_path(flat_mask)\n",
    "\n",
    "\n",
    "model_state_decay_mask(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(1.6668038, dtype=float32)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "State({\n",
       "  'linear1': {\n",
       "    'bias': VariableState(\n",
       "      type=Param,\n",
       "      value=Array([ 0.6486942 , -2.2046442 ,  0.14181204], dtype=float32)\n",
       "    ),\n",
       "    'kernel': VariableState(\n",
       "      type=Param,\n",
       "      value=Array([[-0.50907314,  1.7301297 , -0.11128926],\n",
       "             [ 0.55557084, -1.8881562 ,  0.12145419]], dtype=float32)\n",
       "    )\n",
       "  },\n",
       "  'linear2': {\n",
       "    'bias': VariableState(\n",
       "      type=Param,\n",
       "      value=Array([-0.531473  , -0.71881384, -0.6530258 , -0.664232  ], dtype=float32)\n",
       "    ),\n",
       "    'kernel': VariableState(\n",
       "      type=Param,\n",
       "      value=Array([[ 0.21832745,  0.29528648,  0.26826096,  0.27286443],\n",
       "             [ 0.24969177,  0.33770654,  0.3067986 ,  0.3120634 ],\n",
       "             [-0.01679468, -0.02271469, -0.02063578, -0.0209899 ]],      dtype=float32)\n",
       "    )\n",
       "  }\n",
       "})"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "Array(1.6540173, dtype=float32)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import optax\n",
    "\n",
    "tx = optax.adamw(1e-3, weight_decay=1e-4, mask=model_state_decay_mask(model))\n",
    "state = nnx.Optimizer(model, tx)\n",
    "\n",
    "\n",
    "def loss_fn(model):\n",
    "    return ((model(x) - y) ** 2).mean()\n",
    "\n",
    "\n",
    "loss_fn(model)\n",
    "grads = nnx.grad(loss_fn)(state.model)\n",
    "grads\n",
    "state.update(grads)\n",
    "loss_fn(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{('linear1', 'bias'): 'decay',\n",
       " ('linear1', 'kernel'): 'no_decay',\n",
       " ('linear2', 'bias'): 'decay',\n",
       " ('linear2', 'kernel'): 'no_decay'}"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "State({\n",
       "  'linear1': {\n",
       "    'bias': decay,\n",
       "    'kernel': no_decay\n",
       "  },\n",
       "  'linear2': {\n",
       "    'bias': decay,\n",
       "    'kernel': no_decay\n",
       "  }\n",
       "})"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from flax.traverse_util import path_aware_map\n",
    "\n",
    "\n",
    "def partition_fn(path, x):\n",
    "    # print(path)\n",
    "    if path[0][-1] in (\"bias\"):\n",
    "        return \"decay\"\n",
    "    else:\n",
    "        return \"no_decay\"\n",
    "\n",
    "\n",
    "path_aware_map(\n",
    "    partition_fn,\n",
    "    nnx.state(model).flat_state(),\n",
    ")\n",
    "\n",
    "nnx.State.from_flat_path(\n",
    "    path_aware_map(\n",
    "        partition_fn,\n",
    "        nnx.state(model).flat_state(),\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(1.6796587, dtype=float32)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "State({\n",
       "  'linear1': {\n",
       "    'bias': VariableState(\n",
       "      type=Param,\n",
       "      value=Array([ 0.64893687, -2.2157297 ,  0.14647551], dtype=float32)\n",
       "    ),\n",
       "    'kernel': VariableState(\n",
       "      type=Param,\n",
       "      value=Array([[-0.5092636,  1.7388293, -0.114949 ],\n",
       "             [ 0.5557786, -1.8976502,  0.1254482]], dtype=float32)\n",
       "    )\n",
       "  },\n",
       "  'linear2': {\n",
       "    'bias': VariableState(\n",
       "      type=Param,\n",
       "      value=Array([-0.53430504, -0.7224043 , -0.65456724, -0.666221  ], dtype=float32)\n",
       "    ),\n",
       "    'kernel': VariableState(\n",
       "      type=Param,\n",
       "      value=Array([[ 0.21807964,  0.29485342,  0.26716533,  0.2719219 ],\n",
       "             [ 0.25243354,  0.34130144,  0.30925167,  0.31475753],\n",
       "             [-0.01829538, -0.02473617, -0.02241333, -0.02281238]],      dtype=float32)\n",
       "    )\n",
       "  }\n",
       "})"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from flax.traverse_util import path_aware_map\n",
    "\n",
    "# nnx.state(model).flat_state()\n",
    "\n",
    "\n",
    "def partition_fn(path, x):\n",
    "    # print(path)\n",
    "    if path[0][-1] in (\"bias\"):\n",
    "        return \"decay\"\n",
    "    else:\n",
    "        return \"no_decay\"\n",
    "\n",
    "\n",
    "# param_partitions = flax.core.freeze(\n",
    "#     path_aware_map(\n",
    "#         partition_fn,\n",
    "#         nnx.state(model).flat_state(),\n",
    "#     )\n",
    "# )\n",
    "\n",
    "param_partitions = nnx.State.from_flat_path(\n",
    "    path_aware_map(\n",
    "        partition_fn,\n",
    "        nnx.state(model).flat_state(),\n",
    "    )\n",
    ")\n",
    "\n",
    "\n",
    "def get_optimizer(decay):\n",
    "    return optax.adamw(learning_rate=1e-3, weight_decay=decay)\n",
    "\n",
    "\n",
    "partition_optimizers = {\n",
    "    \"decay\": get_optimizer(0.1),\n",
    "    \"no_decay\": get_optimizer(0.0),\n",
    "}\n",
    "\n",
    "tx = optax.multi_transform(partition_optimizers, param_partitions)\n",
    "\n",
    "state = nnx.Optimizer(model, tx)\n",
    "\n",
    "\n",
    "def loss_fn(model):\n",
    "    return ((model(x) - y) ** 2).mean()\n",
    "\n",
    "\n",
    "loss_fn(model)\n",
    "grads = nnx.grad(loss_fn)(state.model)\n",
    "grads\n",
    "\n",
    "# Would a label function in place of param partitions help?"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "jax",
   "language": "python",
   "name": "python3"
  },
  "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.9.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}