{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4e93353a",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "source": [
    "# GPUs\n",
    "\n",
    "View the graphics card information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a4ef6e6b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:36:58.270913Z",
     "iopub.status.busy": "2023-08-18T19:36:58.270055Z",
     "iopub.status.idle": "2023-08-18T19:37:01.897059Z",
     "shell.execute_reply": "2023-08-18T19:37:01.896067Z"
    },
    "origin_pos": 6,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from d2l import torch as d2l"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63b9bb73",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "Computing Devices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d996a07b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:37:01.901957Z",
     "iopub.status.busy": "2023-08-18T19:37:01.901006Z",
     "iopub.status.idle": "2023-08-18T19:37:01.911076Z",
     "shell.execute_reply": "2023-08-18T19:37:01.909836Z"
    },
    "origin_pos": 12,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(device(type='cpu'),\n",
       " device(type='cuda', index=0),\n",
       " device(type='cuda', index=1))"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def cpu():  \n",
    "    \"\"\"Get the CPU device.\"\"\"\n",
    "    return torch.device('cpu')\n",
    "\n",
    "def gpu(i=0):  \n",
    "    \"\"\"Get a GPU device.\"\"\"\n",
    "    return torch.device(f'cuda:{i}')\n",
    "\n",
    "cpu(), gpu(), gpu(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "128204a0",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "source": [
    "Query the number of available GPUs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b20d4266",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:37:01.915209Z",
     "iopub.status.busy": "2023-08-18T19:37:01.914386Z",
     "iopub.status.idle": "2023-08-18T19:37:01.922363Z",
     "shell.execute_reply": "2023-08-18T19:37:01.921100Z"
    },
    "origin_pos": 15,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def num_gpus():  \n",
    "    \"\"\"Get the number of available GPUs.\"\"\"\n",
    "    return torch.cuda.device_count()\n",
    "\n",
    "num_gpus()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19a66315",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "Define two convenient functions that allow us\n",
    "to run code even if the requested GPUs do not exist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6ac547f6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:37:01.926431Z",
     "iopub.status.busy": "2023-08-18T19:37:01.925574Z",
     "iopub.status.idle": "2023-08-18T19:37:01.935019Z",
     "shell.execute_reply": "2023-08-18T19:37:01.933960Z"
    },
    "origin_pos": 18,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(device(type='cuda', index=0),\n",
       " device(type='cpu'),\n",
       " [device(type='cuda', index=0), device(type='cuda', index=1)])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def try_gpu(i=0):  \n",
    "    \"\"\"Return gpu(i) if exists, otherwise return cpu().\"\"\"\n",
    "    if num_gpus() >= i + 1:\n",
    "        return gpu(i)\n",
    "    return cpu()\n",
    "\n",
    "def try_all_gpus():  \n",
    "    \"\"\"Return all available GPUs, or [cpu(),] if no GPU exists.\"\"\"\n",
    "    return [gpu(i) for i in range(num_gpus())]\n",
    "\n",
    "try_gpu(), try_gpu(10), try_all_gpus()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac9133fc",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "Query the device where the tensor is located"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a3e90ced",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:37:01.939959Z",
     "iopub.status.busy": "2023-08-18T19:37:01.938949Z",
     "iopub.status.idle": "2023-08-18T19:37:01.950067Z",
     "shell.execute_reply": "2023-08-18T19:37:01.949195Z"
    },
    "origin_pos": 24,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cpu')"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.tensor([1, 2, 3])\n",
    "x.device"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5ad6013",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "Store a tensor on the GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "13913886",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:37:01.953772Z",
     "iopub.status.busy": "2023-08-18T19:37:01.953191Z",
     "iopub.status.idle": "2023-08-18T19:37:02.420258Z",
     "shell.execute_reply": "2023-08-18T19:37:02.419290Z"
    },
    "origin_pos": 29,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1., 1., 1.],\n",
       "        [1., 1., 1.]], device='cuda:0')"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X = torch.ones(2, 3, device=try_gpu())\n",
    "X"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a275ca1b",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "source": [
    "Create a random tensor, `Y`, on the second GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6f4c7aff",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:37:02.424924Z",
     "iopub.status.busy": "2023-08-18T19:37:02.424008Z",
     "iopub.status.idle": "2023-08-18T19:37:02.688334Z",
     "shell.execute_reply": "2023-08-18T19:37:02.687371Z"
    },
    "origin_pos": 34,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.0022, 0.5723, 0.2890],\n",
       "        [0.1456, 0.3537, 0.7359]], device='cuda:1')"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Y = torch.rand(2, 3, device=try_gpu(1))\n",
    "Y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ea5b9a7",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "If we want to compute `X + Y`,\n",
    "we need to decide where to perform this operation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3560f0b5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:37:02.693634Z",
     "iopub.status.busy": "2023-08-18T19:37:02.693201Z",
     "iopub.status.idle": "2023-08-18T19:37:02.701839Z",
     "shell.execute_reply": "2023-08-18T19:37:02.701004Z"
    },
    "origin_pos": 39,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1., 1., 1.],\n",
      "        [1., 1., 1.]], device='cuda:0')\n",
      "tensor([[1., 1., 1.],\n",
      "        [1., 1., 1.]], device='cuda:1')\n"
     ]
    }
   ],
   "source": [
    "Z = X.cuda(1)\n",
    "print(X)\n",
    "print(Z)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8fabb477",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "The data (both `Z` and `Y`) are on the same GPU), we can add them up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2cfea6e5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:37:02.707070Z",
     "iopub.status.busy": "2023-08-18T19:37:02.705679Z",
     "iopub.status.idle": "2023-08-18T19:37:02.735588Z",
     "shell.execute_reply": "2023-08-18T19:37:02.734193Z"
    },
    "origin_pos": 43,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1.0022, 1.5723, 1.2890],\n",
       "        [1.1456, 1.3537, 1.7359]], device='cuda:1')"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Y + Z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0450cb7c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:37:02.743585Z",
     "iopub.status.busy": "2023-08-18T19:37:02.743275Z",
     "iopub.status.idle": "2023-08-18T19:37:02.750645Z",
     "shell.execute_reply": "2023-08-18T19:37:02.748215Z"
    },
    "origin_pos": 49,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Z.cuda(1) is Z"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9cceee0",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "Neural Networks and GPUs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "351af69d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:37:02.768539Z",
     "iopub.status.busy": "2023-08-18T19:37:02.767413Z",
     "iopub.status.idle": "2023-08-18T19:37:02.809950Z",
     "shell.execute_reply": "2023-08-18T19:37:02.807298Z"
    },
    "origin_pos": 58,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.7802],\n",
       "        [0.7802]], device='cuda:0', grad_fn=<AddmmBackward0>)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net = nn.Sequential(nn.LazyLinear(1))\n",
    "net = net.to(device=try_gpu())\n",
    "\n",
    "net(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c94ff609",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "source": [
    "Confirm that the model parameters are stored on the same GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6fdbd2c3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:37:02.816317Z",
     "iopub.status.busy": "2023-08-18T19:37:02.815749Z",
     "iopub.status.idle": "2023-08-18T19:37:02.822467Z",
     "shell.execute_reply": "2023-08-18T19:37:02.821657Z"
    },
    "origin_pos": 62,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cuda', index=0)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net[0].weight.data.device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "1283ae3a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-08-18T19:37:02.826029Z",
     "iopub.status.busy": "2023-08-18T19:37:02.825482Z",
     "iopub.status.idle": "2023-08-18T19:37:02.832065Z",
     "shell.execute_reply": "2023-08-18T19:37:02.831156Z"
    },
    "origin_pos": 67,
    "tab": [
     "pytorch"
    ]
   },
   "outputs": [],
   "source": [
    "@d2l.add_to_class(d2l.Trainer)  \n",
    "def __init__(self, max_epochs, num_gpus=0, gradient_clip_val=0):\n",
    "    self.save_hyperparameters()\n",
    "    self.gpus = [d2l.gpu(i) for i in range(min(num_gpus, d2l.num_gpus()))]\n",
    "\n",
    "@d2l.add_to_class(d2l.Trainer)  \n",
    "def prepare_batch(self, batch):\n",
    "    if self.gpus:\n",
    "        batch = [a.to(self.gpus[0]) for a in batch]\n",
    "    return batch\n",
    "\n",
    "@d2l.add_to_class(d2l.Trainer)  \n",
    "def prepare_model(self, model):\n",
    "    model.trainer = self\n",
    "    model.board.xlim = [0, self.max_epochs]\n",
    "    if self.gpus:\n",
    "        model.to(self.gpus[0])\n",
    "    self.model = model"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Slideshow",
  "language_info": {
   "name": "python"
  },
  "required_libs": [],
  "rise": {
   "autolaunch": true,
   "enable_chalkboard": true,
   "overlay": "<div class='my-top-right'><img height=80px src='http://d2l.ai/_static/logo-with-text.png'/></div><div class='my-top-left'></div>",
   "scroll": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}