notebooks/persistence_example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model Persistence Example\n",
    "\n",
    "Model persistence is an important part of deploying and sharing models. This notebook shows the three supported ways for saving and loading models: `torch.save/load`, `pickle` and PyTorch's `state_dict`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.chdir('..')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `torch.save` and `torch.load`\n",
    "\n",
    "The most general-purpose way of saving TextWiser objects is to use [torch.save](https://pytorch.org/tutorials/beginner/saving_loading_models.html#save-load-entire-model). This has the benefit of letting all Torch tensors be handled by PyTorch, while pickling everything else.\n",
    "\n",
    "It is easy, but comes with the caveat that these artifacts are not highly portable: you will need the exact same requirements (apart from PyTorch) in order to load them back in. This method has the added benefit of being able to load TextWiser models that have been trained on a GPU on a CPU environment.\n",
    "\n",
    "To demonstrate, we first create a PyTorch model that uses pooled word vectors, with an SVD in the middle to reduce the dimensionality. The results are then fed into a linear layer which produces a desired output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (0): TextWiser(\n",
       "    (_imp): _Sequential(\n",
       "      (0): _WordEmbeddings(\n",
       "        (model): Embedding(246117, 50)\n",
       "      )\n",
       "      (1): _SVDTransformation()\n",
       "      (2): _PoolTransformation()\n",
       "    )\n",
       "  )\n",
       "  (1): Linear(in_features=2, out_features=1, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from textwiser import TextWiser, Embedding, Transformation, WordOptions, PoolOptions, device\n",
    "\n",
    "docs = ['This is one document.', 'This is a second document.', 'Not the second document!']\n",
    "\n",
    "tw = TextWiser(Embedding.Word(word_option=WordOptions.word2vec, pretrained='en-turian'),\n",
    "                       [Transformation.SVD(n_components=2), Transformation.Pool(pool_option=PoolOptions.mean)], dtype=torch.float32)\n",
    "tw.fit(docs)\n",
    "model = nn.Sequential(tw, nn.Linear(2, 1))\n",
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For demonstration, we can look at the output of the model with some example documents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.2845],\n",
       "        [0.2469],\n",
       "        [0.2470]], grad_fn=<AddmmBackward>)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Get results of the model\n",
    "expected = model(docs)\n",
    "expected"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model can then be persisted using `torch.save` and loaded back in. Note that in the below cell we delete the `model` object between saving and loading, meaning the `model` object after loading is brand new."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.2845],\n",
      "        [0.2469],\n",
      "        [0.2470]], grad_fn=<AddmmBackward>)\n"
     ]
    }
   ],
   "source": [
    "import pickle\n",
    "from tempfile import NamedTemporaryFile\n",
    "\n",
    "with NamedTemporaryFile() as file:\n",
    "    torch.save(model, file.name)\n",
    "    del model\n",
    "    model = torch.load(file.name, map_location=device)\n",
    "    print(model(docs))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected, the output of the loaded object is exactly the same as when it was first created.\n",
    "\n",
    "## Pickle\n",
    "\n",
    "Pickling is the default way of persisting arbitrary Python objects to files. It is easy, but comes with the caveat that these artifacts are not highly portable: you will need the exact same requirements in order to load them back in. Despite that, they remain the default way of persisting Scikit-learn models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.2845],\n",
      "        [0.2469],\n",
      "        [0.2470]], grad_fn=<AddmmBackward>)\n"
     ]
    }
   ],
   "source": [
    "with NamedTemporaryFile() as file:\n",
    "    with open(file.name, 'wb') as fp:\n",
    "        pickle.dump(model, fp)\n",
    "    del model\n",
    "    with open(file.name, 'rb') as fp:\n",
    "        model = pickle.load(fp)\n",
    "    print(model(docs))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As expected, the output of the loaded object is exactly the same as when it was first created.\n",
    "\n",
    "## State dict\n",
    "\n",
    "The preferred way of saving any PyTorch model is to use the state dictionary. An example can be found [here](https://pytorch.org/tutorials/beginner/saving_loading_models.html).\n",
    "\n",
    "TextWiser overloads the state dictionary to also hold the data it requires. For example, an `Embedding.TfIdf` model will store its internal Scikit-learn `TfIdfVectorizer` object inside the state dictionary. Note that these objects will **still get pickled**, meaning that the Scikit-learn version needs to remain the same when saving and loading. However, the PyTorch version **can** change without any issues."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.2845],\n",
      "        [0.2469],\n",
      "        [0.2470]], grad_fn=<AddmmBackward>)\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "with NamedTemporaryFile() as file:\n",
    "    # Save the model\n",
    "    torch.save(model.state_dict(), file.name)\n",
    "    # Get rid of the original model\n",
    "    del tw\n",
    "    del model\n",
    "    # Create the same model\n",
    "    tw = TextWiser(Embedding.Word(word_option=WordOptions.word2vec, pretrained='en-turian'),\n",
    "                   [Transformation.SVD(n_components=2), Transformation.Pool(pool_option=PoolOptions.mean)], dtype=torch.float32)\n",
    "    tw.fit()\n",
    "    model = nn.Sequential(tw, nn.Linear(2, 1))\n",
    "    # Load the model from file\n",
    "    model.load_state_dict(torch.load(file.name))\n",
    "    # Do predictions with the loaded model\n",
    "    predicted = model(docs)\n",
    "    print(predicted)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As shown here, the output is again the same between saving and loading. Again, the model is deleted in between saving and loading.\n",
    "\n",
    "While this approach doesn't fully solve the dependency on a specific environment, it does lessen the blow, and has the added benefit of being compatible with 3rd party experiment management solutions, which assume that your model parameters will be persisted using the state dictionary."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (textwiser)",
   "language": "python",
   "name": "textwiser"
  },
  "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}