examples/04_configure_gpu_memory.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using a single memory pool for Cupy and PyTorch or TensorFlow\n",
    "\n",
    "Requesting memory from a GPU device directly is expensive, so most deep learning libraries will over-allocate, and maintain an internal pool of memory they will keep a hold of, instead of returning it back to the device. This means the libraries don't by default play well together: they all expect to be the single consumer of the GPU memory, so they hog it selfishly. If you use two frameworks together, you can get unexpected out-of-memory errors.\n",
    "\n",
    "Thinc's internal models use cupy for GPU operations, and cupy offers a nice solution for this problem. You can provide cupy with a custom memory allocation function, which allows us to route cupy's memory requests via another library. This avoids the memory problem when you use PyTorch and cupy together, or when you use cupy and Tensorflow together. We don't yet have a similar solution for using PyTorch and Tensorflow together, however.\n",
    "\n",
    "To start with, we call the `require_gpu()` function, which tells Thinc and PyTorch to allocate on GPU."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install \"thinc>=8.0.0\" torch \"tensorflow>=2.0\" "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from thinc.api import require_gpu\n",
    "\n",
    "require_gpu()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then call `use_pytorch_for_gpu_memory()` to set up the allocation strategy. Now when `cupy` tries to request GPU memory, it will do so by asking PyTorch, rather than asking the GPU directly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from thinc.api import use_pytorch_for_gpu_memory\n",
    "\n",
    "use_pytorch_for_gpu_memory()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To test that it's working, we make a little function that allocates an array using cupy, and prints its size, along with the current size of PyTorch's memory pool. Notice the over-allocation: PyTorch grabs a *much* bigger chunk of memory than just our little array. That's why we need to have only one memory pool."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cupy \n",
    "import torch.cuda\n",
    "\n",
    "def allocate_cupy_tensor(size):\n",
    "    array = cupy.zeros((size,), dtype=\"f\")\n",
    "    print(array.size, torch.cuda.max_memory_allocated())\n",
    "    return array\n",
    "allocate_cupy_tensor(16)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also see that even when we free the tensor, the memory isn't immediately released. On the other hand, we don't need to resize the memory pool when we make a second small allocation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow\n",
    "\n",
    "with tensorflow.device('/device:GPU:0'):\n",
    "    arr = allocate_cupy_tensor(1000)\n",
    "    arr = None\n",
    "    arr = allocate_cupy_tensor(1000)\n",
    "    arr = None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we make a huge allocation, we'll have to resize the pool though. Let's make sure the pool resizes properly, and  that memory is freed when the tensors are removed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arr = allocate_cupy_tensor(1000)\n",
    "for _ in range(100):\n",
    "    arr2 = allocate_cupy_tensor(900000)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}