code/17-使用和购买GPU.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79fae070",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "527e2f88",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "94442ac9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "\n",
    "\"\"\"指定cpu, gpu设备\"\"\"\n",
    "torch.device('cpu'), torch.device('cuda'), torch.device('cuda:1')\n",
    "# cpu, gpu0, gpu1\n",
    "\n",
    "\"\"\"查询可用gpu数量\"\"\"\n",
    "torch.cuda.device_count()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "85a5d238",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(device(type='cpu'), device(type='cpu'), [device(type='cpu')])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"定义了两个方便的函数， 这两个函数允许我们在不存在所需所有GPU的情况下运行代码\"\"\"\n",
    "\n",
    "\n",
    "\"\"\"如果存在，则返回gpu(i)，否则返回cpu()\"\"\"\n",
    "def try_gpu(i=0):\n",
    "# 不输入参数则默认i = 0 \n",
    "    if torch.cuda.device_count() >= i + 1:\n",
    "        return torch.device(f'cuda:{i}')\n",
    "     # 如果当前可用gpu的总数大于等于i+1，则返回第i个gpu（从0计数）\t\t\n",
    "    return torch.device('cpu')\n",
    "\t# 否则证明当前没有更多可用gpu，则返回cpu\n",
    "\n",
    "\n",
    "\"\"\"返回所有可用的GPU，如果没有GPU，则返回[cpu(),]\"\"\"\n",
    "def try_all_gpus():\n",
    "    devices = [torch.device(f'cuda:{i}')\n",
    "             for i in range(torch.cuda.device_count())]\n",
    "    # 所有可用gpu设备序号组成的列表devices\n",
    "    \n",
    "    return devices if devices else [torch.device('cpu')]\n",
    "\t# 如果列表devices不为空则证明此时有可用的gpu，则返回可用gpu序号列表；否则证明没有可用gpu，则返回cpu\n",
    "\n",
    "try_gpu(), try_gpu(10), try_all_gpus()\n",
    "# 测试函数功能\n",
    "# try_gpu():检测是否有第i=0号gpu\n",
    "# try_gpu(10):检测是否有第i=10号gpu\n",
    "# try_all_gpus():返回所有可用gpu序号列表，如果没有gpu则返回cpu\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "295a6578",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cpu')"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"我们可以查询张量所在的设备。 默认情况下，张量是在CPU上创建的。\"\"\"\n",
    "x = torch.tensor([1, 2, 3])\n",
    "x.device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "255462a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"我们在第一个gpu上创建张量变量X\"\"\" \n",
    "X = torch.ones(2, 3, device=try_gpu())\n",
    "\n",
    "\"\"\"假设你至少有两个GPU，下面的代码将在第二个GPU上创建一个随机张量\"\"\"\n",
    "Y = torch.rand(2, 3, device=try_gpu(1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f72119d",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"将gpu(0)中的X复制到gpu(1)中的Z\"\"\"\n",
    "Z = X.cuda(1)\n",
    "print(X)\n",
    "print(Z)\n",
    "\n",
    "\n",
    "\"\"\"现在数据在同一个GPU上（Z和Y都在），我们可以将它们相加。\"\"\"\n",
    "Y + Z\n",
    "\n",
    "\n",
    "\"\"\"如果变量Z已经存在于第i个GPU上，再调用Z.cuda(i)只会返回Z并不会复制并分配新内存\"\"\"\n",
    "Z.cuda(1) is Z\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "e9861ae0",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "类似地，神经网络模型可以指定设备。 下面的代码将模型参数放在GPU上。\n",
    "\"\"\"\n",
    "\n",
    "net = nn.Sequential(nn.Linear(3, 1))\n",
    "net = net.to(device=try_gpu())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "0b17de5d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cpu')"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net(X)\n",
    "\n",
    "net[0].weight.data.device\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26d4de02",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}