Merge pull request #756 from ZhouLong0/Modified

Added 'getting Pytorch to run on Apple Arm GPU' guide

Author: mrdbourke

00_pytorch_fundamentals.ipynb

@@ -3501,6 +3501,86 @@
     "Knowing the number of GPUs PyTorch has access to is helpful incase you wanted to run a specific process on one GPU and another process on another (PyTorch also has features to let you run a process across *all* GPUs)."
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "\n",
+    "\n",
+    "### 2.1 Getting PyTorch to run on the ARM GPUs\n",
+    "\n",
+    "In order to run PyTorch on the Apple's M1/M2 GPUs you can use the [`torch.backends.mps`](https://pytorch.org/docs/stable/notes/mps.html) package.\n",
+    "\n",
+    "Be sure that the versions of the MacOS and Pytorch are updated\n",
+    "\n",
+    "You can test if PyTorch has access to a GPU using `torch.backends.mps.is_available()`\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "True"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Check for ARM GPU\n",
+    "import torch\n",
+    "torch.backends.mps.is_available()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'mps'"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Set device type\n",
+    "device = \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n",
+    "device"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "As before, if the above output `\"mps\"` it means we can set all of our PyTorch code to use the available Apple Arm GPU"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if torch.cuda.is_available():\n",
+    "    device = 'cuda'\n",
+    "elif torch.backends.mps.is_available():\n",
+    "    device = 'mps'\n",
+    "else:\n",
+    "    device = 'cpu'"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {
@@ -3524,7 +3604,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 74,
+   "execution_count": 9,
    "metadata": {
     "colab": {
      "base_uri": "https://localhost:8080/"
@@ -3543,10 +3623,10 @@
     {
      "data": {
       "text/plain": [
-       "tensor([1, 2, 3], device='cuda:0')"
+       "tensor([1, 2, 3], device='mps:0')"
       ]
      },
-     "execution_count": 74,
+     "execution_count": 9,
      "metadata": {},
      "output_type": "execute_result"
     }