test_ivy_integration.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-09-16 07:16:51.442856: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
     ]
    }
   ],
   "source": [
    "from thinc.api import IvyWrapper, Adam, CategoricalCrossentropy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ivy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ml_datasets\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training size=54000, dev size=10000\n"
     ]
    }
   ],
   "source": [
    "(train_X, train_Y), (dev_X, dev_Y) = ml_datasets.mnist()\n",
    "print(f\"Training size={len(train_X)}, dev size={len(dev_X)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(data, model, optimizer, n_iter, batch_size):\n",
    "    (train_X, train_Y), (dev_X, dev_Y) = data\n",
    "    indices = model.ops.xp.arange(train_X.shape[0], dtype=\"i\")\n",
    "    for i in range(n_iter):\n",
    "        batches = model.ops.multibatch(batch_size, train_X, train_Y, shuffle=True)\n",
    "        for X, Y in tqdm(batches, leave=False):\n",
    "            Yh, backprop = model.begin_update(X)\n",
    "            backprop(Yh - Y)\n",
    "            model.finish_update(optimizer)\n",
    "        # Evaluate and print progress\n",
    "        correct = 0\n",
    "        total = 0\n",
    "        for X, Y in model.ops.multibatch(batch_size, dev_X, dev_Y):\n",
    "            Yh = model.predict(X)\n",
    "            correct += (Yh.argmax(axis=1) == Y.argmax(axis=1)).sum()\n",
    "            total += Yh.shape[0]\n",
    "        score = correct / total\n",
    "        print(f\" {i} {float(score):.3f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "width = 32\n",
    "nO = 10\n",
    "nI = 784\n",
    "dropout = 0.2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Regressor(ivy.Module):\n",
    "    def __init__(self, width, nO, nI, dropout):\n",
    "        self.nO = nO\n",
    "        self.nI = nI\n",
    "        self.dropout = dropout\n",
    "        self.width = width\n",
    "        super().__init__()\n",
    "\n",
    "    def _build(self, *args, **kwargs):\n",
    "        self.dropout1 = ivy.Dropout(self.dropout)\n",
    "        self.dropout2 = ivy.Dropout(self.dropout)\n",
    "        self.fc1 = ivy.Linear(self.nI, self.width)\n",
    "        self.fc2 = ivy.Linear(self.width, self.nO)\n",
    "\n",
    "    def _forward(self, x):\n",
    "        x = ivy.relu(x)\n",
    "        x = self.dropout1(x)\n",
    "        x = self.fc1(x)\n",
    "        x = ivy.relu(x)\n",
    "        x = self.dropout2(x)\n",
    "        x = self.fc2(x)\n",
    "        output = ivy.log_softmax(x, axis=1)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = IvyWrapper(\n",
    "    Regressor(width, nO, nI, dropout),\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                 \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 0 0.489\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                 \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 1 0.542\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                 \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 2 0.609\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                 \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 3 0.678\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                 \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 4 0.674\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                 \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 5 0.667\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                 \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 6 0.683\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                 \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 7 0.685\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                 \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 8 0.674\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                 \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 9 0.672\n"
     ]
    }
   ],
   "source": [
    "data = ml_datasets.mnist()\n",
    "optimizer = Adam(0.001)\n",
    "model.initialize(X=train_X[:5], Y=train_Y[:5])\n",
    "train_model(data, model, optimizer, n_iter=10, batch_size=128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "thinc",
   "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.10.0"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}