06_01_TPU.ipynb

{"nbformat":4,"nbformat_minor":0,"metadata":{"colab":{"name":"06_01_TPU.ipynb","provenance":[{"file_id":"11xQTsGWASrUIQE3usoaDRzNG_rofchsL","timestamp":1615395403638}],"collapsed_sections":[],"authorship_tag":"ABX9TyP8HB9m+OyKsmi3NwPdmlAP"},"kernelspec":{"name":"python3","display_name":"Python 3"},"accelerator":"TPU"},"cells":[{"cell_type":"markdown","metadata":{"id":"Y2fkP_np3Zpl"},"source":["<!--BOOK_INFORMATION-->\n","<img align=\"left\" style=\"width:80px;height:98px;padding-right:20px;\" src=\"https://raw.githubusercontent.com/joe-papa/pytorch-book/main/files/pytorch-book-cover.jpg\">\n","\n","This notebook contains an excerpt from the [PyTorch Pocket Reference](http://pytorchbook.com) book by [Joe Papa](http://joepapa.ai); content is available [on GitHub](https://github.com/joe-papa/pytorch-book)."]},{"cell_type":"markdown","metadata":{"id":"Ew8mdVxDHU9t"},"source":["[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/joe-papa/pytorch-book/blob/main/06_01_TPU.ipynb)"]},{"cell_type":"markdown","metadata":{"id":"aXkszegKMMMn"},"source":["# Chapter 6 - PyTorch Acceleration & Optimization"]},{"cell_type":"markdown","metadata":{"id":"QbzMwCbOLF4T"},"source":["## PyTorch on a Tensor Processing Unit (TPU)\n","\n"]},{"cell_type":"code","metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"HWv6NA5YKpLU","executionInfo":{"status":"ok","timestamp":1615397429320,"user_tz":300,"elapsed":49557,"user":{"displayName":"Joe Papa","photoUrl":"","userId":"00487850786587503652"}},"outputId":"32da966c-870f-48b9-8b44-650e6b380eb1"},"source":["!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n","!python pytorch-xla-env-setup.py --version \"nightly\""],"execution_count":null,"outputs":[{"output_type":"stream","text":["  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n","                                 Dload  Upload   Total   Spent    Left  Speed\n","\r  0     0    0     0    0     0      0      0 --:--:-- --:--:-- --:--:--     0\r100  5116  100  5116    0     0  85266      0 --:--:-- --:--:-- --:--:-- 86711\n","Updating... This may take around 2 minutes.\n","Updating TPU runtime to pytorch-nightly ...\n","Uninstalling torch-1.9.0a0+git6aa5148:\n","  Successfully uninstalled torch-1.9.0a0+git6aa5148\n","Uninstalling torchvision-0.9.0a0+551193b:\n","  Successfully uninstalled torchvision-0.9.0a0+551193b\n","Copying gs://tpu-pytorch/wheels/torch-nightly-cp37-cp37m-linux_x86_64.whl...\n","\\ [1 files][127.2 MiB/127.2 MiB]                                                \n","Operation completed over 1 objects/127.2 MiB.                                    \n","Copying gs://tpu-pytorch/wheels/torch_xla-nightly-cp37-cp37m-linux_x86_64.whl...\n","\\ [1 files][138.2 MiB/138.2 MiB]                                                \n","Operation completed over 1 objects/138.2 MiB.                                    \n","Copying gs://tpu-pytorch/wheels/torchvision-nightly-cp37-cp37m-linux_x86_64.whl...\n","/ [1 files][  4.5 MiB/  4.5 MiB]                                                \n","Operation completed over 1 objects/4.5 MiB.                                      \n","Processing ./torch-nightly-cp37-cp37m-linux_x86_64.whl\n","Done updating TPU runtime\n","Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch==nightly) (3.7.4.3)\n","\u001b[31mERROR: fastai 1.0.61 requires torchvision, which is not installed.\u001b[0m\n","\u001b[31mERROR: torchtext 0.9.0 has requirement torch==1.8.0, but you'll have torch 1.9.0a0+git6aa5148 which is incompatible.\u001b[0m\n","Installing collected packages: torch\n","Successfully installed torch-1.9.0a0+git6aa5148\n","Processing ./torch_xla-nightly-cp37-cp37m-linux_x86_64.whl\n","Installing collected packages: torch-xla\n","  Found existing installation: torch-xla 1.9+9fc2c18\n","    Uninstalling torch-xla-1.9+9fc2c18:\n","      Successfully uninstalled torch-xla-1.9+9fc2c18\n","Successfully installed torch-xla-1.9+9fc2c18\n","Processing ./torchvision-nightly-cp37-cp37m-linux_x86_64.whl\n","Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from torchvision==nightly) (1.19.5)\n","Requirement already satisfied: torch in /usr/local/lib/python3.7/dist-packages (from torchvision==nightly) (1.9.0a0+git6aa5148)\n","Requirement already satisfied: pillow>=4.1.1 in /usr/local/lib/python3.7/dist-packages (from torchvision==nightly) (7.0.0)\n","Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch->torchvision==nightly) (3.7.4.3)\n","Installing collected packages: torchvision\n","Successfully installed torchvision-0.9.0a0+551193b\n","Reading package lists... Done\n","Building dependency tree       \n","Reading state information... Done\n","libomp5 is already the newest version (5.0.1-1).\n","0 upgraded, 0 newly installed, 0 to remove and 29 not upgraded.\n"],"name":"stdout"}]},{"cell_type":"code","metadata":{"id":"rgMYPb-DLP3i"},"source":["import torch_xla.core.xla_model as xm\n","\n","device = xm.xla_device()"],"execution_count":null,"outputs":[]},{"cell_type":"code","metadata":{"id":"9oeufUi3En9I"},"source":["# Dummy values to get code to run in the next cells\n","import torch\n","from torch import nn, optim\n","from torch.utils.data import DataLoader\n","\n","n_epochs = 1\n","model = nn.Linear(10,10)\n","dataset = [(torch.rand(10,requires_grad=True),torch.rand(10,requires_grad=True))]*20\n","\n","trainloader = DataLoader(dataset)\n","criterion = nn.MSELoss()\n","optimizer = optim.SGD(model.parameters(),lr=0.001)"],"execution_count":null,"outputs":[]},{"cell_type":"code","metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"1mYqkKMKI1VW","executionInfo":{"status":"ok","timestamp":1615396602941,"user_tz":300,"elapsed":206,"user":{"displayName":"Joe Papa","photoUrl":"","userId":"00487850786587503652"}},"outputId":"0b2c8a0c-4cbf-425d-8a21-13e194558e26"},"source":["device = torch.device(\"cuda\" if\n","  torch.cuda.is_available() else \"cpu\")\n","print(device)"],"execution_count":null,"outputs":[{"output_type":"stream","text":["cpu\n"],"name":"stdout"}]},{"cell_type":"code","metadata":{"id":"a-uuMJOSC695","colab":{"base_uri":"https://localhost:8080/"},"executionInfo":{"status":"ok","timestamp":1615396483286,"user_tz":300,"elapsed":519,"user":{"displayName":"Joe Papa","photoUrl":"","userId":"00487850786587503652"}},"outputId":"b6274b8e-8307-41ab-ad59-8859a384be84"},"source":["model.to(device)\n","for epoch in range(n_epochs):\n","  for data in trainloader:\n","    input, labels = data\n","    input = input.to(device)\n","    labels = labels.to(device)\n","    optimizer.zero_grad()\n","    \n","    output = model(input)\n","    \n","    loss = criterion(input, labels)\n","    loss.backward()\n","    optimizer.step()\n","    \n","print(output.device)\n","# out: xla:1"],"execution_count":null,"outputs":[{"output_type":"stream","text":["xla:1\n"],"name":"stdout"}]},{"cell_type":"code","metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"JF3fKHvBHllf","executionInfo":{"status":"ok","timestamp":1615397021957,"user_tz":300,"elapsed":228,"user":{"displayName":"Joe Papa","photoUrl":"","userId":"00487850786587503652"}},"outputId":"9ace2d95-90ed-430e-8cf2-5b0f675c3d00"},"source":["xm.xla_device()"],"execution_count":null,"outputs":[{"output_type":"execute_result","data":{"text/plain":["device(type='xla', index=1)"]},"metadata":{"tags":[]},"execution_count":24}]},{"cell_type":"code","metadata":{"id":"9MrjYL8GKlCi"},"source":[""],"execution_count":null,"outputs":[]}]}