Copy_of_00_intro_to_thinc.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "file_extension": ".py",
    "kernelspec": {
      "display_name": "Python 3.7.2 64-bit ('.env': venv)",
      "language": "python",
      "name": "python37"
    },
    "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.7.2"
    },
    "mimetype": "text/x-python",
    "name": "python",
    "npconvert_exporter": "python",
    "pygments_lexer": "ipython3",
    "version": 3,
    "colab": {
      "name": "Copy of 00_intro_to_thinc.ipynb",
      "provenance": [],
      "include_colab_link": true
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/jonbaer/googlecolab/blob/master/Copy_of_00_intro_to_thinc.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lbQzd_LzIAxN",
        "colab_type": "text"
      },
      "source": [
        "# Intro to Thinc for beginners: defining a simple model and config & wrapping PyTorch, TensorFlow and MXNet\n",
        "\n",
        "This example shows how to get started with Thinc, using the \"hello world\" of neural network models: recognizing handwritten digits from the [MNIST dataset](http://yann.lecun.com/exdb/mnist/). For comparison, here's the same model implemented in other frameworks: [PyTorch version](https://github.com/pytorch/examples/blob/master/mnist/main.py), [TensorFlow version](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/examples/tutorials/mnist/mnist.py). In this notebook, we'll walk through **creating and training the model**, using **config files**, registering **custom functions** and **wrapping models** defined in PyTorch, TensorFlow and MXNet. This tutorial is aimed at beginners, but it assumes basic knowledge of machine learning concepts and terminology."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2ujceyQMIAxP",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install \"thinc>=8.0.0a0\" ml_datasets \"tqdm>=4.41\""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bZ2Cb29sIAxT",
        "colab_type": "text"
      },
      "source": [
        "First, let's use Thinc's `prefer_gpu` helper to make sure we're performing operations **on GPU if available**. The function should be called right after importing Thinc, and it returns a boolean indicating whether the GPU has been activated."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FPHeVDzGIAxU",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from thinc.api import prefer_gpu\n",
        "prefer_gpu()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9eykCWm3IAxX",
        "colab_type": "text"
      },
      "source": [
        "We’ve prepared a separate package [`ml-datasets`](https://github.com/explosion/ml-datasets) with loaders for some common datasets, including MNIST. So we can set up the data as follows:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "i9lbhJk3IAxZ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import ml_datasets\n",
        "(train_X, train_Y), (dev_X, dev_Y) = ml_datasets.mnist()\n",
        "print(f\"Training size={len(train_X)}, dev size={len(dev_X)}\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vCjXmGuIIAxc",
        "colab_type": "text"
      },
      "source": [
        "Now let’s define a model with two **Relu-activated hidden layers**, followed by a **softmax-activated output layer**. We’ll also add **dropout** after the two hidden layers, to help the model generalize better. The `chain` combinator is like `Sequential` in PyTorch or Keras: it combines a list of layers together with a feed-forward relationship."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NoAWwsLgIAxd",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from thinc.api import chain, Relu, Softmax\n",
        " \n",
        "n_hidden = 32\n",
        "dropout = 0.2\n",
        "\n",
        "model = chain(\n",
        "    Relu(nO=n_hidden, dropout=dropout), \n",
        "    Relu(nO=n_hidden, dropout=dropout), \n",
        "    Softmax()\n",
        ")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2VORt68BIAxh",
        "colab_type": "text"
      },
      "source": [
        "After creating the model, we can call the `Model.initialize` method, passing in a small batch of input data `X` and a small batch of output data `Y`. This allows Thinc to **infer the missing dimensions**: when we defined the model, we didn’t tell it the input size `nI` or the output size `nO`. When passing in the data, make sure it is on the right device by calling `model.ops.asarray` which will e.g. transform the arrays to `cupy` when running on GPU."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5TMRPhxnIAxi",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# making sure the data is on the right device\n",
        "train_X = model.ops.asarray(train_X)\n",
        "train_Y = model.ops.asarray(train_Y)\n",
        "dev_X = model.ops.asarray(dev_X)\n",
        "dev_Y = model.ops.asarray(dev_Y)\n",
        "\n",
        "model.initialize(X=train_X[:5], Y=train_Y[:5])\n",
        "nI = model.get_dim(\"nI\")\n",
        "nO = model.get_dim(\"nO\")\n",
        "print(f\"Initialized model with input dimension nI={nI} and output dimension nO={nO}\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0BxIhS3IIAxl",
        "colab_type": "text"
      },
      "source": [
        "Next we need to create an **optimizer**, and make several passes over the data, randomly selecting paired batches of the inputs and labels each time. While some machine learning libraries provide a single `.fit()` method to train a model all at once, Thinc puts you in charge of **shuffling and batching your data**, with the help of a few handy utility methods. `model.ops.xp` is an instance of either `numpy` or `cupy`, depending on whether you run the code on CPU or GPU."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uh8X9e6VIAxo",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from thinc.api import Adam, fix_random_seed\n",
        "from tqdm.notebook import tqdm\n",
        "\n",
        "fix_random_seed(0)\n",
        "optimizer = Adam(0.001)\n",
        "batch_size = 128\n",
        "print(\"Measuring performance across iterations:\")\n",
        "\n",
        "for i in range(10):\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}\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AGnDsCVRIAxs",
        "colab_type": "text"
      },
      "source": [
        "Let's wrap the training code in a function, so we can reuse it later:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fY7sdA1hIAxs",
        "colab_type": "code",
        "colab": {}
      },
      "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}\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RK_UmqGiIAxv",
        "colab_type": "text"
      },
      "source": [
        "### Operator overloading for more concise model definitions\n",
        "\n",
        "Thinc allows you to **overload operators** and bind arbitrary functions to Python operators like `+`, `*`, but also `>>` or `@`. The `Model.define_operators` contextmanager takes a dict of operators mapped to functions – typically combinators like `chain`. The operators are only valid for the `with` block. This lets us define the model like this:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9Qb-fHCEIAxw",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from thinc.api import Model, chain, Relu, Softmax\n",
        " \n",
        "n_hidden = 32\n",
        "dropout = 0.2\n",
        "\n",
        "with Model.define_operators({\">>\": chain}):\n",
        "    model = Relu(nO=n_hidden, dropout=dropout) >> Relu(nO=n_hidden, dropout=dropout) >> Softmax()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jN7iICg2IAxz",
        "colab_type": "text"
      },
      "source": [
        "If your model definitions are very complex, mapping combinators to operators can help you keep the code readable and concise. You can find more examples of model definitions with overloaded operators [in the docs](https://thinc.ai/docs). (Also note that you don't _have to_ use this syntax!)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BJlvJG55IAx0",
        "colab_type": "text"
      },
      "source": [
        "---\n",
        "\n",
        "## Using config files\n",
        "\n",
        "Configuration is a huge problem for machine learning code, because you may want to expose almost any detail of any function as a hyperparameter. The setting you want to expose might be arbitrarily far down in your call stack. Default values also become hard to change without breaking backwards compatibility.\n",
        "\n",
        "To solve this problem, Thinc provides a config system that lets you easily describe **arbitrary trees of objects**. The objects can be created via function calls you register using a simple decorator syntax. The config can include values like hyperparameters or training settings (whatever you need), or references to functions and the values of their arguments. Thinc will then construct the config **bottom-up** – so you can define one function with its arguments, and then pass the return value into another function.\n",
        "\n",
        "> 💡 You can keep the config as a string in your Python script, or save it to a file like `config.cfg`. To load a config from a string, you can use `Config.from_str`. To load from a file, you can use `Config.from_disk`. The following examples all use strings so we can include them in the notebook."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Vrdwqj2PIAx1",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from thinc.api import Config, registry\n",
        "\n",
        "EXAMPLE_CONFIG1 = \"\"\"\n",
        "[hyper_params]\n",
        "learn_rate = 0.001\n",
        "\n",
        "[optimizer]\n",
        "@optimizers = \"Adam.v1\"\n",
        "learn_rate = ${hyper_params:learn_rate}\n",
        "\"\"\"\n",
        "\n",
        "config1 = Config().from_str(EXAMPLE_CONFIG1)\n",
        "config1"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p7FA-GVUIAx3",
        "colab_type": "text"
      },
      "source": [
        "When you open the config with `Config.from_str`, Thinc will parse it as a dict and fill in the references to values defined in other sections. For example, `${hyper_params:learn_rate}` is substituted with `0.001`. \n",
        "\n",
        "Keys starting with `@` are references to **registered functions**. For example, `@optimizers = \"Adam.v1\"` refers to the function registered under the name `\"Adam.v1\"`, a function creating an Adam optimizer. The function takes one argument, the `learn_rate`. Calling `registry.make_from_config` will resolve the config and create the functions it defines."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9tfIKDp6IAx4",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "loaded_config1 = registry.make_from_config(config1)\n",
        "loaded_config1"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NtI-YJlyIAx7",
        "colab_type": "text"
      },
      "source": [
        "If function arguments are missing or have incompatible types, Thinc will raise an error and tell you what's wrong. Configs can also define **nested blocks** using the `.` notation. In this example, `optimizer.learn_rate` defines the `learn_rate` argument of the `optimizer` block. Instead of a float, the learning rate can also be a generator – for instance, a linear warm-up rate:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tO8FED1KIAx8",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "EXAMPLE_CONFIG2 = \"\"\"\n",
        "[optimizer]\n",
        "@optimizers = \"Adam.v1\"\n",
        "\n",
        "[optimizer.learn_rate]\n",
        "@schedules = \"warmup_linear.v1\"\n",
        "initial_rate = 2e-5\n",
        "warmup_steps = 1000\n",
        "total_steps = 10000\n",
        "\"\"\"\n",
        "\n",
        "config2 = Config().from_str(EXAMPLE_CONFIG2)\n",
        "config2"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1rtgEEwnIAyA",
        "colab_type": "text"
      },
      "source": [
        "Calling `registry.make_from_config` will now construct the objects bottom-up: first, it will create the schedule with the given arguments. Next, it will create the optimizer and pass in the schedule as the `learn_rate` argument."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cDacJSCPIAyB",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "loaded_config2 = registry.make_from_config(config2)\n",
        "loaded_config2"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5OmG7SmtIAyE",
        "colab_type": "text"
      },
      "source": [
        "This gives you a loaded optimizer using the settings defined in the config, which you can then use in your script. How you set up your config and what you do with the result is **entirely up to you**. Thinc just gives you a dictionary of objects back and makes no assumptions about what they _\"mean\"_. This means that you can also choose the names of the config sections – the only thing that needs to stay consistent are the names of the function arguments."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ltsmMgwsIAyF",
        "colab_type": "text"
      },
      "source": [
        "### Configuring the MNIST model\n",
        "\n",
        "Here's a config describing the model we defined above. The values in the `hyper_params` section can be referenced in other sections to keep them consistent. The `*` is used for **positional arguments** – in this case, the arguments to the `chain` function, two Relu layers and one softmax layer."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "35lVVsJEIAyG",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "CONFIG = \"\"\"\n",
        "[hyper_params]\n",
        "n_hidden = 32\n",
        "dropout = 0.2\n",
        "learn_rate = 0.001\n",
        "\n",
        "[model]\n",
        "@layers = \"chain.v1\"\n",
        "\n",
        "[model.*.relu1]\n",
        "@layers = \"Relu.v1\"\n",
        "nO = ${hyper_params:n_hidden}\n",
        "dropout = ${hyper_params:dropout}\n",
        "\n",
        "[model.*.relu2]\n",
        "@layers = \"Relu.v1\"\n",
        "nO = ${hyper_params:n_hidden}\n",
        "dropout = ${hyper_params:dropout}\n",
        "\n",
        "[model.*.softmax]\n",
        "@layers = \"Softmax.v1\"\n",
        "\n",
        "[optimizer]\n",
        "@optimizers = \"Adam.v1\"\n",
        "learn_rate = ${hyper_params:learn_rate}\n",
        "\n",
        "[training]\n",
        "n_iter = 10\n",
        "batch_size = 128\n",
        "\"\"\"\n",
        "\n",
        "config = Config().from_str(CONFIG)\n",
        "config"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eFmJBnRwIAyL",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "loaded_config = registry.make_from_config(config)\n",
        "loaded_config"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u4lTrYajIAyN",
        "colab_type": "text"
      },
      "source": [
        "When you call `registry.make_from_config`, Thinc will first create the three layers using the specified arguments populated by the hyperparameters. It will then pass the return values (the layer objects) to `chain`. It will also create an optimizer. All other values, like the training config, will be passed through as a regular dict. Your training code can now look like this:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pu39xNXzIAyO",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "model = loaded_config[\"model\"]\n",
        "optimizer = loaded_config[\"optimizer\"]\n",
        "n_iter = loaded_config[\"training\"][\"n_iter\"]\n",
        "batch_size = loaded_config[\"training\"][\"batch_size\"]\n",
        "\n",
        "model.initialize(X=train_X[:5], Y=train_Y[:5])\n",
        "train_model(((train_X, train_Y), (dev_X, dev_Y)), model, optimizer, n_iter, batch_size)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C86HB291IAyR",
        "colab_type": "text"
      },
      "source": [
        "If you want to change a hyperparamter or experiment with a different optimizer, all you need to change is the config. For each experiment you run, you can save a config and you'll be able to reproduce it later."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iwWO4euWIAyS",
        "colab_type": "text"
      },
      "source": [
        "---\n",
        "\n",
        "## Programming via config vs. registering custom functions\n",
        "\n",
        "The config system is very powerful and lets you define complex relationships, including model definitions with levels of nested layers. However, it's not always a good idea to program entirely in your config – this just replaces one problem (messy and hard to maintain code) with another one (messy and hard to maintain configs). So ultimately, it's about finding the **best possible trade-off**.\n",
        "\n",
        "If you've written a layer or model definition you're happy with, you can use Thinc's function registry to register it and assign it a string name. Your function can take any arguments that can later be defined in the config. Adding **type hints** ensures that config settings will be **parsed and validated** before they're passed into the function, so you don't end up with incompatible settings and confusing failures later on. Here's the MNIST model, defined as a custom layer:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5zOGp2kmIAyT",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import thinc\n",
        "\n",
        "@thinc.registry.layers(\"MNIST.v1\")\n",
        "def create_mnist(nO: int, dropout: float):\n",
        "    return chain(\n",
        "        Relu(nO, dropout=dropout), \n",
        "        Relu(nO, dropout=dropout), \n",
        "        Softmax()\n",
        "    )"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M69VZVXoIAyW",
        "colab_type": "text"
      },
      "source": [
        "In the config, we can now refer to it by name and set its arguments. This makes the config maintainable and compact, while still allowing you to change and record the hyperparameters."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_lBPSGn0IAyX",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "CONFIG2 = \"\"\"\n",
        "[model]\n",
        "@layers = \"MNIST.v1\"\n",
        "nO = 32\n",
        "dropout = 0.2\n",
        "\n",
        "[optimizer]\n",
        "@optimizers = \"Adam.v1\"\n",
        "learn_rate = 0.001\n",
        "\n",
        "[training]\n",
        "n_iter = 10\n",
        "batch_size = 128\n",
        "\"\"\"\n",
        "\n",
        "config = Config().from_str(CONFIG2)\n",
        "config"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mbZhwl-xIAyZ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "loaded_config = registry.make_from_config(config)\n",
        "loaded_config"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G1-nzl71IAyc",
        "colab_type": "text"
      },
      "source": [
        "If you don't want to hard-code the dataset being used, you can also wrap it in a registry function. This lets you refer to it by name in the config, and makes it easy to swap it out. In your config, you can then load the data in its own section, or as a subsection of `training`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X2uytSaVIAyd",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "@thinc.registry.datasets(\"mnist_data.v1\")\n",
        "def mnist():\n",
        "    return ml_datasets.mnist()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RaZsrSU9IAyf",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "CONFIG3 = \"\"\"\n",
        "[model]\n",
        "@layers = \"MNIST.v1\"\n",
        "nO = 32\n",
        "dropout = 0.2\n",
        "\n",
        "[optimizer]\n",
        "@optimizers = \"Adam.v1\"\n",
        "learn_rate = 0.001\n",
        "\n",
        "[training]\n",
        "n_iter = 10\n",
        "batch_size = 128\n",
        "\n",
        "[training.data]\n",
        "@datasets = \"mnist_data.v1\"\n",
        "\"\"\"\n",
        "\n",
        "config = Config().from_str(CONFIG3)\n",
        "loaded_config = registry.make_from_config(config)\n",
        "loaded_config"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SYCJpF0CIAyi",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "model = loaded_config[\"model\"]\n",
        "optimizer = loaded_config[\"optimizer\"]\n",
        "n_iter = loaded_config[\"training\"][\"n_iter\"]\n",
        "batch_size = loaded_config[\"training\"][\"batch_size\"]\n",
        "(train_X, train_Y), (dev_X, dev_Y) = loaded_config[\"training\"][\"data\"]\n",
        "\n",
        "model.initialize(X=train_X[:5], Y=train_Y[:5])\n",
        "train_model(((train_X, train_Y), (dev_X, dev_Y)), model, optimizer, n_iter, batch_size)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4Hlr2V0WIAyk",
        "colab_type": "text"
      },
      "source": [
        "---\n",
        "\n",
        "## Wrapping TensorFlow, PyTorch and MXNet models\n",
        "\n",
        "The previous example showed how to define the model directly in Thinc, which is pretty straightforward. But you can also define your model using a **machine learning library of your choice** and wrap it as a Thinc model. This gives your layers a unified interface so you can easily mix and match them, and also lets you take advantage of the config system and type hints. Thinc currently ships with built-in wrappers for [PyTorch](https://pytorch.org), [TensorFlow](https://tensorflow.org) and [MXNet](https://mxnet.apache.org/)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yBWbTHvRIAyl",
        "colab_type": "text"
      },
      "source": [
        "### Wrapping TensorFlow models\n",
        "\n",
        "Here's the same model definition in TensorFlow: a `Sequential` layer (equivalent of Thinc's `chain`) with two Relu layers and dropout, and an output layer with a softmax activation. Thinc's `TensorFlowWrapper` wraps the model and turns it into a regular Thinc `Model`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nucOM4KPIAym",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install \"tensorflow>2.0\""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OaeYq1mLIAyp",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from tensorflow.keras.layers import Dense, Dropout\n",
        "from tensorflow.keras.models import Sequential\n",
        "from thinc.api import TensorFlowWrapper, Adam\n",
        "\n",
        "width = 32\n",
        "nO = 10\n",
        "nI = 784\n",
        "dropout = 0.2\n",
        "\n",
        "tf_model = Sequential()\n",
        "tf_model.add(Dense(width, activation=\"relu\", input_shape=(nI,)))\n",
        "tf_model.add(Dropout(dropout))\n",
        "tf_model.add(Dense(width, activation=\"relu\", input_shape=(nI,)))\n",
        "tf_model.add(Dropout(dropout))\n",
        "tf_model.add(Dense(nO, activation=\"softmax\"))\n",
        "\n",
        "wrapped_tf_model = TensorFlowWrapper(tf_model)\n",
        "wrapped_tf_model"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NxgidAefIAyu",
        "colab_type": "text"
      },
      "source": [
        "You can now use the same training code to train the model:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LLpuKP4sIAyv",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "data = ml_datasets.mnist()\n",
        "optimizer = Adam(0.001)\n",
        "wrapped_tf_model.initialize(X=train_X[:5], Y=train_Y[:5])\n",
        "train_model(data, wrapped_tf_model, optimizer, n_iter=10, batch_size=128)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "teZvq9UPIAyx",
        "colab_type": "text"
      },
      "source": [
        "### Wrapping PyTorch models\n",
        "\n",
        "Here's the PyTorch version. Thinc's `PyTorchWrapper` wraps the model and turns it into a regular Thinc `Model`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qN2TPtL3IAyy",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install torch"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_DUELnsgIAy0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torch\n",
        "import torch.nn\n",
        "import torch.nn.functional as F\n",
        "from thinc.api import PyTorchWrapper, Adam\n",
        "\n",
        "\n",
        "width = 32\n",
        "nO = 10\n",
        "nI = 784\n",
        "dropout = 0.2\n",
        "\n",
        "\n",
        "class PyTorchModel(torch.nn.Module):\n",
        "    def __init__(self, width, nO, nI, dropout):\n",
        "        super(PyTorchModel, self).__init__()\n",
        "        self.dropout1 = torch.nn.Dropout2d(dropout)\n",
        "        self.dropout2 = torch.nn.Dropout2d(dropout)\n",
        "        self.fc1 = torch.nn.Linear(nI, width)\n",
        "        self.fc2 = torch.nn.Linear(width, nO)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = F.relu(x)\n",
        "        x = self.dropout1(x)\n",
        "        x = self.fc1(x)\n",
        "        x = F.relu(x)\n",
        "        x = self.dropout2(x)\n",
        "        x = self.fc2(x)\n",
        "        output = F.log_softmax(x, dim=1)\n",
        "        return output\n",
        "    \n",
        "wrapped_pt_model = PyTorchWrapper(PyTorchModel(width, nO, nI, dropout))\n",
        "wrapped_pt_model"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QPZHj5HlIAy3",
        "colab_type": "text"
      },
      "source": [
        "You can now use the same training code to train the model:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9JnzvCq8IAy4",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "data = ml_datasets.mnist()\n",
        "optimizer = Adam(0.001)\n",
        "wrapped_pt_model.initialize(X=train_X[:5], Y=train_Y[:5])\n",
        "train_model(data, wrapped_pt_model, optimizer, n_iter=10, batch_size=128)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oJn_sHYVIAy7",
        "colab_type": "text"
      },
      "source": [
        "### Wrapping MXNet models\n",
        "\n",
        "Here's the MXNet version. Thinc's `MXNetWrapper` wraps the model and turns it into a regular Thinc `Model`.\n",
        "\n",
        "MXNet doesn't provide a `Softmax` layer but a `.softmax()` operation/method for prediction and it integrates an internal softmax during training. So to be able to integrate it with the rest of the components, you combine it with a `Softmax()` Thinc layer using the `chain` combinator. Make sure you `initialize()` the MXNet model *and* the Thinc model."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CUCCMZUXIAy9",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install \"mxnet>=1.5.1,<1.6.0\""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qjgAl4doIAy_",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from mxnet.gluon.nn import Dense, Sequential, Dropout\n",
        "from thinc.api import MXNetWrapper, chain, Softmax\n",
        "\n",
        "width = 32\n",
        "nO = 10\n",
        "nI = 784\n",
        "dropout = 0.2\n",
        "\n",
        "mx_model = Sequential()\n",
        "mx_model.add(Dense(width, activation=\"relu\"))\n",
        "mx_model.add(Dropout(dropout))\n",
        "mx_model.add(Dense(width, activation=\"relu\"))\n",
        "mx_model.add(Dropout(dropout))\n",
        "mx_model.add(Dense(nO))\n",
        "mx_model.initialize()\n",
        "wrapped_mx_model = chain(MXNetWrapper(mx_model), Softmax())\n",
        "wrapped_mx_model"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RTnZeRybIAzB",
        "colab_type": "text"
      },
      "source": [
        "And train it the same way:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z4Rgeg56IAzC",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "data = ml_datasets.mnist()\n",
        "optimizer = Adam(0.001)\n",
        "wrapped_mx_model.initialize(X=train_X[:5], Y=train_Y[:5])\n",
        "train_model(data, wrapped_mx_model, optimizer, n_iter=10, batch_size=128)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LZaaaEICIAzF",
        "colab_type": "text"
      },
      "source": [
        "---\n",
        "\n",
        "## Documentation and resources\n",
        "\n",
        "- <kbd>USAGE</kbd> [Configuration files](https://thinc.ai/docs/usage-config)\n",
        "- <kbd>USAGE</kbd> [Defining and using models](https://thinc.ai/docs/usage-models)\n",
        "- <kbd>USAGE</kbd> [Using Thinc with PyTorch, TensorFlow & MXNet](https://thinc.ai/docs/usage-frameworks)\n",
        "- <kbd>API</kbd> [Available layers and combinators](https://thinc.ai/docs/api-layers)\n",
        "- <kbd>API</kbd> [`Config` and `registry`](https://thinc.ai/docs/api-config)\n",
        "- <kbd>API</kbd> [`Model` class](https://thinc.ai/docs/api-model)"
      ]
    }
  ]
}