gpt2-nano training

nbs/01_tensor.ipynb

@@ -33,7 +33,8 @@
     "from lovely_numpy import lovely\n",
     "\n",
     "import tidygrad.ops as ops\n",
-    "import tidygrad.tensor_helpers as helpers"
+    "import tidygrad.tensor_helpers as helpers\n",
+    "import traceback"
    ]
   },
   {
@@ -67,21 +68,54 @@
     "class Tensor:\n",
     "    pass\n",
     "\n",
+    "def simplify_trace(trace):\n",
+    "    return ' -> '.join(f'{frame.name} at {frame.filename}:{frame.lineno}' for frame in trace if '/python' not in frame.filename)\n",
+    "\n",
+    "alloc_log = {}\n",
+    "\n",
     "class Tensor:\n",
-    "    # op = \"L\"\n",
     "    name: str = \"\"\n",
     "\n",
     "    def __init__(self, data, name=None, op=None, eps=1e-8, requires_grad=False):\n",
     "        global _num_tensors\n",
     "        _num_tensors += 1\n",
-    "        self.data = np.asarray(data, dtype=np.float64) # , dtype=np.float32\n",
     "\n",
-    "        self.grad = (np.zeros_like(self.data, dtype=np.float64) if requires_grad else None)\n",
+    "        trace = traceback.extract_stack()\n",
+    "        simplified_trace = simplify_trace(trace)\n",
+    "        alloc_log[id(self)] = simplified_trace\n",
+    "        \n",
+    "        # Increment allocation count\n",
+    "\n",
+    "        # if _num_tensors > 620:\n",
+    "        #     raise Exception(\"Too many tensors\")\n",
+    "\n",
+    "        self.data = np.asarray(data) # , dtype=np.float32\n",
+    "        if self.data.dtype == np.float64:\n",
+    "            self.data = self.data.astype(np.float32)\n",
+    "\n",
+    "        self.grad = (np.zeros_like(self.data, dtype=np.float32) if requires_grad else None)\n",
     "        self.eps = eps\n",
     "        self.op = op or ops.Load(name=name)\n",
     "        self.name = name or self.op.name\n",
-    "        self.requires_grad = requires_grad\n",
+    "        self._requires_grad = requires_grad\n",
     "\n",
+    "    def __del__(self):\n",
+    "        # print(f\"Tensor {self.name} deleted\")\n",
+    "        del alloc_log[id(self)]\n",
+    "        global _num_tensors\n",
+    "        _num_tensors -= 1\n",
+    "\n",
+    "    @property\n",
+    "    def requires_grad(self):\n",
+    "        return self._requires_grad\n",
+    "\n",
+    "    @requires_grad.setter\n",
+    "    def requires_grad(self, requires_grad):\n",
+    "        if requires_grad and self.grad is None:\n",
+    "            self.grad = np.zeros_like(self.data)\n",
+    "        \n",
+    "        self._requires_grad = requires_grad\n",
+    "    \n",
     "    def __repr__(self):\n",
     "        value_str = f\"v={lovely(self.data)}\"\n",
     "        grad_str = f\"∇={lovely(self.grad)}\" if self.grad is not None else \"\"\n",
@@ -90,7 +124,7 @@
     "        return f'Tensor{list(self.data.shape)}(name=\"{self.name}\" op={type(self.op).__name__}{parents}):\\n    {value_str}\\n    {grad_str}'\n",
     "\n",
     "    def accum_grad(self, grad):\n",
-    "        if not self.requires_grad:\n",
+    "        if not self._requires_grad:\n",
     "            return\n",
     "\n",
     "        if self.grad is None:\n",
@@ -230,9 +264,11 @@
     "        for n in nodes[::-1]:\n",
     "            if hasattr(n.op, \"backward\"):\n",
     "                n.op.backward()\n",
+    "                n.op = None\n",
+    "\n",
     "\n",
     "    def zero_grad(self):\n",
-    "        assert self.requires_grad, \"Cannot zero grad on non-differentiable tensor\"\n",
+    "        assert self._requires_grad, \"Cannot zero grad on non-differentiable tensor\"\n",
     "        self.grad.fill(0)"
    ]
   }

nbs/02_func.ipynb

@@ -215,6 +215,7 @@
     "        target = Tensor(target)\n",
     "    sm = softmax(logits)\n",
     "    loss = -target * sm.log()\n",
+    "\n",
     "    if reduction == \"mean\":\n",
     "        return loss.mean(axis=-1, keepdims=True)\n",
     "    if reduction == \"sum\":\n",

nbs/02_ops.conv.ipynb

@@ -1,5 +1,14 @@
 {
  "cells": [
+  {
+   "cell_type": "raw",
+   "metadata": {},
+   "source": [
+    "---\n",
+    "skip_exec: true\n",
+    "---"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,

nbs/10_utils.grad_check.ipynb

@@ -190,7 +190,7 @@
     "\n",
     "loss.backward()\n",
     "\n",
-    "grad_check(NN, (x, y), (w1, b1, w2))"
+    "# grad_check(NN, (x, y), (w1, b1, w2))"
    ]
   }
  ],