deploy: 955a9c0

Author: jxbz

.doctrees/environment.pickle

@@ -80,30 +80,33 @@
   },
   {
    "cell_type": "markdown",
-   "id": "a7af8a3d-77dc-4007-a6de-03ab617bf3fa",
+   "id": "d7d5fc30",
    "metadata": {},
    "source": [
-    "Next, we choose our error measure. Error measures allow us to both compute the loss of the model, and also to compute the derivative of the loss with respect to model outputs. For simplicity we will just use squared error."
+    "Next, we set up a loss function and create a jitted function for both evaluating the loss and also returning its gradient."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 3,
-   "id": "a7ea38a1-2684-437f-88fc-cb1f2a44133c",
+   "id": "8b719f14",
    "metadata": {},
    "outputs": [],
    "source": [
-    "from modula.error import SquareError\n",
+    "def mse(w, inputs, targets):\n",
+    "    outputs = mlp(inputs, w)\n",
+    "    loss = ((outputs-targets) ** 2).mean()\n",
+    "    return loss\n",
     "\n",
-    "error = SquareError()"
+    "mse_and_grad = jax.jit(jax.value_and_grad(mse))"
    ]
   },
   {
    "cell_type": "markdown",
    "id": "1c4b8252-b3f0-4d16-9b48-9d8d582c1abe",
    "metadata": {},
    "source": [
-    "Finally we are ready to train our model. The method `mlp.backward` takes as input the weights, activations and the gradient of the error. It returns the gradient of the loss with respect to both the model weights and the inputs. The method `mlp.dualize` takes in the gradient of the weights and solves for the vector of unit modular norm that maximizes the linearized improvement in loss."
+    "Finally we are ready to train our model. We will apply the method `mlp.dualize` to the gradient of the loss to solve for the vector of unit modular norm that maximizes the linearized improvement in loss."
    ]
   },
   {
@@ -116,16 +119,16 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Step   0 \t Loss 0.976274\n",
-      "Step 100 \t Loss 0.001989\n",
-      "Step 200 \t Loss 0.001537\n",
-      "Step 300 \t Loss 0.001194\n",
-      "Step 400 \t Loss 0.000885\n",
-      "Step 500 \t Loss 0.000627\n",
-      "Step 600 \t Loss 0.000420\n",
-      "Step 700 \t Loss 0.000255\n",
-      "Step 800 \t Loss 0.000134\n",
-      "Step 900 \t Loss 0.000053\n"
+      "Step   0 \t Loss 0.976154\n",
+      "Step 100 \t Loss 0.001773\n",
+      "Step 200 \t Loss 0.001371\n",
+      "Step 300 \t Loss 0.001002\n",
+      "Step 400 \t Loss 0.000696\n",
+      "Step 500 \t Loss 0.000453\n",
+      "Step 600 \t Loss 0.000282\n",
+      "Step 700 \t Loss 0.000152\n",
+      "Step 800 \t Loss 0.000061\n",
+      "Step 900 \t Loss 0.000011\n"
      ]
     }
    ],
@@ -137,17 +140,9 @@
     "w = mlp.initialize(key)\n",
     "\n",
     "for step in range(steps):\n",
-    "    # compute outputs and activations\n",
-    "    outputs, activations = mlp(inputs, w)\n",
-    "    \n",
-    "    # compute loss\n",
-    "    loss = error(outputs, targets)\n",
-    "    \n",
-    "    # compute error gradient\n",
-    "    error_grad = error.grad(outputs, targets)\n",
-    "    \n",
-    "    # compute gradient of weights\n",
-    "    grad_w, _ = mlp.backward(w, activations, error_grad)\n",
+    "\n",
+    "    # compute loss and gradient of weights\n",
+    "    loss, grad_w = mse_and_grad(w, inputs, targets)\n",
     "    \n",
     "    # dualize gradient\n",
     "    d_w = mlp.dualize(grad_w)\n",
@@ -179,7 +174,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.16"
+   "version": "3.12.8"
   }
  },
  "nbformat": 4,