Skip to content

Commit

Permalink
Add neural style transfer notebook.
Browse files Browse the repository at this point in the history
  • Loading branch information
@dlibenzi
Ailing Zhang authored and dlibenzi committed Dec 3, 2019
1 parent ec87d32 commit cd46571
Showing 1 changed file with 245 additions and 0 deletions.
245 changes: 245 additions & 0 deletions contrib/colab/style_transfer_inference-xrt-1-15.ipynb
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,245 @@
{
"nbformat": 4,
"nbformat_minor": 0,
"metadata": {
"colab": {
"name": "PyTorch/TPU style transfer",
"provenance": [],
"collapsed_sections": [],
"machine_shape": "hm"
},
"kernelspec": {
"name": "python3",
"display_name": "Python 3"
},
"accelerator": "TPU"
},
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "h1mYGqkc1kqv",
"colab_type": "text"
},
"source": [
"## This demo shows how to run existing PyTorch code on Colab TPUs. \n",
"\n",
"- [WARNING]: this notebook requires access to your google drive.\n",
"- This demo is inference only. For training part, please take a look at [PyTorch examples on TPU](https://github.com/pytorch-tpu/examples)\n",
"- This demo only uses 1 TPU core due to small inference workload. To distribute workload across multiple cores, please use https://github.com/pytorch/xla/blob/master/API_GUIDE.md#running-on-multiple-xla-devices-with-multiprocessing"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "xOp9jBEumdvC",
"colab_type": "text"
},
"source": [
"<h3> &nbsp;&nbsp;Use Colab Cloud TPU&nbsp;&nbsp; <a href=\"https://cloud.google.com/tpu/\"><img valign=\"middle\" src=\"https://raw.githubusercontent.com/GoogleCloudPlatform/tensorflow-without-a-phd/master/tensorflow-rl-pong/images/tpu-hexagon.png\" width=\"50\"></a></h3>\n",
"\n",
"* On the main menu, click Runtime and select **Change runtime type**. Set \"TPU\" as the hardware accelerator.\n",
"* The cell below makes sure you have access to a TPU on Colab.\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "YofXQrnxmf5r",
"colab_type": "text"
},
"source": [
"### [RUNME] Install Colab TPU compatible PyTorch/TPU wheels and dependencies\n",
"This may take up to ~2 minutes"
]
},
{
"cell_type": "code",
"metadata": {
"id": "sPJVqAKyml5W",
"colab_type": "code",
"colab": {}
},
"source": [
"import collections\n",
"from datetime import datetime, timedelta\n",
"import os\n",
"import requests\n",
"import threading\n",
"\n",
"_VersionConfig = collections.namedtuple('_VersionConfig', 'wheels,server')\n",
"VERSION = \"xrt==1.15.0\" #@param [\"xrt==1.15.0\", \"torch_xla==nightly\"]\n",
"CONFIG = {\n",
" 'xrt==1.15.0': _VersionConfig('1.15', '1.15.0'),\n",
" 'torch_xla==nightly': _VersionConfig('nightly', 'XRT-dev{}'.format(\n",
" (datetime.today() - timedelta(1)).strftime('%Y%m%d'))),\n",
"}[VERSION]\n",
"DIST_BUCKET = 'gs://tpu-pytorch/wheels'\n",
"TORCH_WHEEL = 'torch-{}-cp36-cp36m-linux_x86_64.whl'.format(CONFIG.wheels)\n",
"TORCH_XLA_WHEEL = 'torch_xla-{}-cp36-cp36m-linux_x86_64.whl'.format(CONFIG.wheels)\n",
"TORCHVISION_WHEEL = 'torchvision-{}-cp36-cp36m-linux_x86_64.whl'.format(CONFIG.wheels)\n",
"\n",
"# Update TPU XRT version\n",
"def update_server_xrt():\n",
" print('Updating server-side XRT to {} ...'.format(CONFIG.server))\n",
" url = 'http://{TPU_ADDRESS}:8475/requestversion/{XRT_VERSION}'.format(\n",
" TPU_ADDRESS=os.environ['COLAB_TPU_ADDR'].split(':')[0],\n",
" XRT_VERSION=CONFIG.server,\n",
" )\n",
" print('Done updating server-side XRT: {}'.format(requests.post(url)))\n",
"\n",
"update = threading.Thread(target=update_server_xrt)\n",
"update.start()\n",
"\n",
"# Install Colab TPU compat PyTorch/TPU wheels and dependencies\n",
"!pip uninstall -y torch torchvision\n",
"!gsutil cp \"$DIST_BUCKET/$TORCH_WHEEL\" .\n",
"!gsutil cp \"$DIST_BUCKET/$TORCH_XLA_WHEEL\" .\n",
"!gsutil cp \"$DIST_BUCKET/$TORCHVISION_WHEEL\" .\n",
"!pip install \"$TORCH_WHEEL\"\n",
"!pip install \"$TORCH_XLA_WHEEL\"\n",
"!pip install \"$TORCHVISION_WHEEL\"\n",
"!sudo apt-get install libomp5\n",
"update.join()"
],
"execution_count": 0,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "Hx4YVNHametU",
"colab_type": "code",
"colab": {}
},
"source": [
"import os\n",
"assert os.environ['COLAB_TPU_ADDR'], 'Make sure to select TPU from Edit > Notebook settings > Hardware accelerator'\n",
"from google.colab.patches import cv2_imshow\n",
"import cv2\n",
"import sys\n",
"from google.colab import drive\n",
"drive.mount('/content/gdrive')"
],
"execution_count": 0,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "aeSKvzuS5LY8",
"colab_type": "code",
"colab": {}
},
"source": [
"# Setup repo in google drive\n",
"REPO_DIR='/content/gdrive/My Drive/demo'\n",
"%mkdir -p \"$REPO_DIR\"\n",
"%cd \"$REPO_DIR\" \n",
"%rm -rf examples\n",
"!git clone https://github.com/pytorch/examples.git \n",
"%cd \"$REPO_DIR/examples/fast_neural_style\"\n",
"\n",
"# Download pretrained weights for styles\n",
"!python download_saved_models.py\n",
"%cd \"$REPO_DIR/examples/fast_neural_style/neural_style\""
],
"execution_count": 0,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "LNLegt4jLkRD",
"colab_type": "code",
"colab": {}
},
"source": [
"import argparse\n",
"import os\n",
"import sys\n",
"import time\n",
"import re\n",
"\n",
"import numpy as np\n",
"import torch\n",
"from torch.optim import Adam\n",
"from torch.utils.data import DataLoader\n",
"from torchvision import datasets\n",
"from torchvision import transforms\n",
"import torch_xla\n",
"import torch_xla.core.xla_model as xm\n",
"import torch_xla.debug.metrics as met\n",
"import torch_xla.distributed.parallel_loader as pl\n",
"import torch_xla.distributed.xla_multiprocessing as xmp\n",
"import torch_xla.utils.utils as xu\n",
"import utils\n",
"from transformer_net import TransformerNet\n",
"from vgg import Vgg16"
],
"execution_count": 0,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "EozMXwIV9iOJ",
"colab_type": "code",
"colab": {},
"cellView": "form"
},
"source": [
"#@markdown ### Enter the url to image(right click -> copy image address):\n",
"content_image_url = 'https://i.pinimg.com/originals/cd/a5/1b/cda51b70758b5302de57225b7c52a202.jpg' #@param {type:\"string\"}\n",
"\n",
"content_image = 'content.jpg'\n",
"!wget -O \"$content_image\" \"$content_image_url\"\n",
"model = '../saved_models/rain_princess.pth'\n",
"RESULT_IMAGE = '/tmp/result.jpg'\n",
"!rm -f \"$RESULT_IMAGE\"\n",
"img = cv2.imread(content_image, cv2.IMREAD_UNCHANGED)\n",
"cv2_imshow(img)"
],
"execution_count": 0,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "yeN4MXqffFEk",
"colab_type": "code",
"colab": {}
},
"source": [
"device = xm.xla_device()\n",
"\n",
"content_image = utils.load_image(content_image, scale=None)\n",
"content_transform = transforms.Compose([\n",
" transforms.ToTensor(),\n",
" transforms.Lambda(lambda x: x.mul(255))\n",
" ])\n",
"content_image = content_transform(content_image)\n",
"content_image = content_image.unsqueeze(0).to(device)\n",
"\n",
"with torch.no_grad():\n",
" style_model = TransformerNet()\n",
" state_dict = torch.load(model)\n",
" # remove saved deprecated running_* keys in InstanceNorm from the checkpoint\n",
" for k in list(state_dict.keys()):\n",
" if re.search(r'in\\d+\\.running_(mean|var)$', k):\n",
" del state_dict[k]\n",
" style_model.load_state_dict(state_dict)\n",
" style_model.to(device)\n",
" print('Running on device: ', content_image.device)\n",
" output = style_model(content_image).cpu()\n",
"\n",
"utils.save_image(RESULT_IMAGE, output[0])\n",
"img = cv2.imread(RESULT_IMAGE, cv2.IMREAD_UNCHANGED)\n",
"\n",
"cv2_imshow(img)\n"
],
"execution_count": 0,
"outputs": []
}
]
}

0 comments on commit cd46571

Please sign in to comment.