site/en/hub/tutorials/s3gan_generation_with_tf_hub.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BhN1AplL0Hpv"
      },
      "source": [
        "##### Copyright 2019 The TensorFlow Hub Authors.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LMgeG2swVVi6"
      },
      "outputs": [],
      "source": [
        "# Copyright 2019 The TensorFlow Hub Authors. All Rights Reserved.\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     http://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "# =============================================================================="
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AqBuuwrIxlGs"
      },
      "source": [
        "# Generating Images with Little Data Using S3GAN\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MfBg1C5NB3X0"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/hub/tutorials/s3gan_generation_with_tf_hub\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/hub/tutorials/s3gan_generation_with_tf_hub.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/hub/tutorials/s3gan_generation_with_tf_hub.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/docs/site/en/hub/tutorials/s3gan_generation_with_tf_hub.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://tfhub.dev/s?q=s3gan\"><img src=\"https://www.tensorflow.org/images/hub_logo_32px.png\" />See TF Hub models</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p5AWAusyySDA"
      },
      "source": [
        "This notebook is a demo of Generative Adversarial Networks trained on ImageNet with as little as 2.5% labeled data using self- and semi-supervised learning techniques. Both generator and discriminator models are available on [TF Hub](https://tfhub.dev/s?publisher=google&q=compare_gan).\n",
        "\n",
        "For more information about the models and the training procedure see our [blogpost](https://ai.googleblog.com/2019/03/reducing-need-for-labeled-data-in.html) and the [paper](https://arxiv.org/abs/1903.02271) [1].\n",
        "The code for training these models is available on [GitHub](https://github.com/google/compare_gan).\n",
        "\n",
        "To get started, connect to a runtime and follow these steps:\n",
        "\n",
        "1. (Optional) Select a model in the second code cell below.\n",
        "2. Click **Runtime > Run all** to run each cell in order.\n",
        "  * Afterwards, the interactive visualizations should update automatically when you modify the settings using the sliders and dropdown menus.\n",
        "\n",
        "Note: if you run into any issues, you can try restarting the runtime and rerunning all cells from scratch by clicking **Runtime > Restart and run all...**.\n",
        "\n",
        "[1] Mario Lucic\\*, Michael Tschannen\\*, Marvin Ritter\\*, Xiaohua Zhai, Olivier\n",
        "    Bachem, Sylvain Gelly, [High-Fidelity Image Generation With Fewer Labels](https://arxiv.org/abs/1903.02271), ICML 2019."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_m5jsOM9kXWP"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NhlMa_tHs0_W"
      },
      "outputs": [],
      "source": [
        "# @title Imports and utility functions\n",
        "import os\n",
        "\n",
        "import IPython\n",
        "from IPython.display import display\n",
        "import numpy as np\n",
        "import PIL.Image\n",
        "import pandas as pd\n",
        "import six\n",
        "\n",
        "import tensorflow.compat.v1 as tf\n",
        "tf.disable_v2_behavior()\n",
        "\n",
        "import tensorflow_hub as hub\n",
        "\n",
        "def imgrid(imarray, cols=8, pad=1):\n",
        "  pad = int(pad)\n",
        "  assert pad >= 0\n",
        "  cols = int(cols)\n",
        "  assert cols >= 1\n",
        "  N, H, W, C = imarray.shape\n",
        "  rows = int(np.ceil(N / float(cols)))\n",
        "  batch_pad = rows * cols - N\n",
        "  assert batch_pad >= 0\n",
        "  post_pad = [batch_pad, pad, pad, 0]\n",
        "  pad_arg = [[0, p] for p in post_pad]\n",
        "  imarray = np.pad(imarray, pad_arg, 'constant')\n",
        "  H += pad\n",
        "  W += pad\n",
        "  grid = (imarray\n",
        "          .reshape(rows, cols, H, W, C)\n",
        "          .transpose(0, 2, 1, 3, 4)\n",
        "          .reshape(rows*H, cols*W, C))\n",
        "  return grid[:-pad, :-pad]\n",
        "\n",
        "\n",
        "def imshow(a, format='png', jpeg_fallback=True):\n",
        "  a = np.asarray(a, dtype=np.uint8)\n",
        "  if six.PY3:\n",
        "    str_file = six.BytesIO()\n",
        "  else:\n",
        "    str_file = six.StringIO()\n",
        "  PIL.Image.fromarray(a).save(str_file, format)\n",
        "  png_data = str_file.getvalue()\n",
        "  try:\n",
        "    disp = display(IPython.display.Image(png_data))\n",
        "  except IOError:\n",
        "    if jpeg_fallback and format != 'jpeg':\n",
        "      print ('Warning: image was too large to display in format \"{}\"; '\n",
        "             'trying jpeg instead.').format(format)\n",
        "      return imshow(a, format='jpeg')\n",
        "    else:\n",
        "      raise\n",
        "  return disp\n",
        "\n",
        "\n",
        "class Generator(object):\n",
        "\n",
        "  def __init__(self, module_spec):\n",
        "    self._module_spec = module_spec\n",
        "    self._sess = None\n",
        "    self._graph = tf.Graph()\n",
        "    self._load_model()\n",
        "\n",
        "  @property\n",
        "  def z_dim(self):\n",
        "    return self._z.shape[-1].value\n",
        "\n",
        "  @property\n",
        "  def conditional(self):\n",
        "    return self._labels is not None\n",
        "\n",
        "  def _load_model(self):\n",
        "    with self._graph.as_default():\n",
        "      self._generator = hub.Module(self._module_spec, name=\"gen_module\",\n",
        "                                   tags={\"gen\", \"bsNone\"})\n",
        "      input_info = self._generator.get_input_info_dict()\n",
        "      inputs = {k: tf.placeholder(v.dtype, v.get_shape().as_list(), k)\n",
        "                for k, v in self._generator.get_input_info_dict().items()}\n",
        "      self._samples = self._generator(inputs=inputs, as_dict=True)[\"generated\"]\n",
        "      print(\"Inputs:\", inputs)\n",
        "      print(\"Outputs:\", self._samples)\n",
        "      self._z = inputs[\"z\"]\n",
        "      self._labels = inputs.get(\"labels\", None)\n",
        "\n",
        "  def _init_session(self):\n",
        "    if self._sess is None:\n",
        "      self._sess = tf.Session(graph=self._graph)\n",
        "      self._sess.run(tf.global_variables_initializer())\n",
        "\n",
        "  def get_noise(self, num_samples, seed=None):\n",
        "    if np.isscalar(seed):\n",
        "      np.random.seed(seed)\n",
        "      return np.random.normal(size=[num_samples, self.z_dim])\n",
        "    z = np.empty(shape=(len(seed), self.z_dim), dtype=np.float32)\n",
        "    for i, s in enumerate(seed):\n",
        "      np.random.seed(s)\n",
        "      z[i] = np.random.normal(size=[self.z_dim])\n",
        "    return z\n",
        "\n",
        "  def get_samples(self, z, labels=None):\n",
        "    with self._graph.as_default():\n",
        "      self._init_session()\n",
        "      feed_dict = {self._z: z}\n",
        "      if self.conditional:\n",
        "        assert labels is not None\n",
        "        assert labels.shape[0] == z.shape[0]\n",
        "        feed_dict[self._labels] = labels\n",
        "      samples = self._sess.run(self._samples, feed_dict=feed_dict)\n",
        "      return np.uint8(np.clip(256 * samples, 0, 255))\n",
        "\n",
        "\n",
        "class Discriminator(object):\n",
        "\n",
        "  def __init__(self, module_spec):\n",
        "    self._module_spec = module_spec\n",
        "    self._sess = None\n",
        "    self._graph = tf.Graph()\n",
        "    self._load_model()\n",
        "\n",
        "  @property\n",
        "  def conditional(self):\n",
        "    return \"labels\" in self._inputs\n",
        "\n",
        "  @property\n",
        "  def image_shape(self):\n",
        "    return self._inputs[\"images\"].shape.as_list()[1:]\n",
        "\n",
        "  def _load_model(self):\n",
        "    with self._graph.as_default():\n",
        "      self._discriminator = hub.Module(self._module_spec, name=\"disc_module\",\n",
        "                                       tags={\"disc\", \"bsNone\"})\n",
        "      input_info = self._discriminator.get_input_info_dict()\n",
        "      self._inputs = {k: tf.placeholder(v.dtype, v.get_shape().as_list(), k)\n",
        "                      for k, v in input_info.items()}\n",
        "      self._outputs = self._discriminator(inputs=self._inputs, as_dict=True)\n",
        "      print(\"Inputs:\", self._inputs)\n",
        "      print(\"Outputs:\", self._outputs)\n",
        "\n",
        "  def _init_session(self):\n",
        "    if self._sess is None:\n",
        "      self._sess = tf.Session(graph=self._graph)\n",
        "      self._sess.run(tf.global_variables_initializer())\n",
        "\n",
        "  def predict(self, images, labels=None):\n",
        "    with self._graph.as_default():\n",
        "      self._init_session()\n",
        "      feed_dict = {self._inputs[\"images\"]: images}\n",
        "      if \"labels\" in self._inputs:\n",
        "        assert labels is not None\n",
        "        assert labels.shape[0] == images.shape[0]\n",
        "        feed_dict[self._inputs[\"labels\"]] = labels\n",
        "      return self._sess.run(self._outputs, feed_dict=feed_dict)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "msTFS1UPkugr"
      },
      "source": [
        "## Select a model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-hBEi9IFdoI-"
      },
      "outputs": [],
      "source": [
        "# @title Select a model { run: \"auto\" }\n",
        "\n",
        "model_name = \"S3GAN 128x128 20% labels (FID 6.9, IS 98.1)\"  # @param [\"S3GAN 256x256 10% labels (FID 8.8, IS 130.7)\", \"S3GAN 128x128 2.5% labels (FID 12.6, IS 48.7)\", \"S3GAN 128x128 5% labels (FID 8.4, IS 74.0)\", \"S3GAN 128x128 10% labels (FID 7.6, IS 90.3)\", \"S3GAN 128x128 20% labels (FID 6.9, IS 98.1)\"]\n",
        "models = {\n",
        "    \"S3GAN 256x256 10% labels\": \"https://tfhub.dev/google/compare_gan/s3gan_10_256x256/1\",\n",
        "    \"S3GAN 128x128 2.5% labels\": \"https://tfhub.dev/google/compare_gan/s3gan_2_5_128x128/1\",\n",
        "    \"S3GAN 128x128 5% labels\": \"https://tfhub.dev/google/compare_gan/s3gan_5_128x128/1\",\n",
        "    \"S3GAN 128x128 10% labels\": \"https://tfhub.dev/google/compare_gan/s3gan_10_128x128/1\",\n",
        "    \"S3GAN 128x128 20% labels\": \"https://tfhub.dev/google/compare_gan/s3gan_20_128x128/1\",\n",
        "}\n",
        "\n",
        "module_spec = models[model_name.split(\" (\")[0]]\n",
        "print(\"Module spec:\", module_spec)\n",
        "\n",
        "tf.reset_default_graph()\n",
        "print(\"Loading model...\")\n",
        "sampler = Generator(module_spec)\n",
        "print(\"Model loaded.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ePQuAme_kxLj"
      },
      "source": [
        "## Sample"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "kGgTXtFYq_FV"
      },
      "outputs": [],
      "source": [
        "# @title Sampling { run: \"auto\" }\n",
        "\n",
        "num_rows = 2  # @param {type: \"slider\", min:1, max:16}\n",
        "num_cols = 3  # @param {type: \"slider\", min:1, max:16}\n",
        "noise_seed = 23  # @param {type:\"slider\", min:0, max:100, step:1}\n",
        "label_str = \"980) volcano\"  # @param [\"-1) Random\", \"0) tench, Tinca tinca\", \"1) goldfish, Carassius auratus\", \"2) great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias\", \"3) tiger shark, Galeocerdo cuvieri\", \"4) hammerhead, hammerhead shark\", \"5) electric ray, crampfish, numbfish, torpedo\", \"6) stingray\", \"7) cock\", \"8) hen\", \"9) ostrich, Struthio camelus\", \"10) brambling, Fringilla montifringilla\", \"11) goldfinch, Carduelis carduelis\", \"12) house finch, linnet, Carpodacus mexicanus\", \"13) junco, snowbird\", \"14) indigo bunting, indigo finch, indigo bird, Passerina cyanea\", \"15) robin, American robin, Turdus migratorius\", \"16) bulbul\", \"17) jay\", \"18) magpie\", \"19) chickadee\", \"20) water ouzel, dipper\", \"21) kite\", \"22) bald eagle, American eagle, Haliaeetus leucocephalus\", \"23) vulture\", \"24) great grey owl, great gray owl, Strix nebulosa\", \"25) European fire salamander, Salamandra salamandra\", \"980) volcano\"]\n",
        "\n",
        "num_samples = num_rows * num_cols\n",
        "z = sampler.get_noise(num_samples, seed=noise_seed)\n",
        "\n",
        "label = int(label_str.split(')')[0])\n",
        "if label == -1:\n",
        "  labels = np.random.randint(0, num_classes, size=(num_samples))\n",
        "else:\n",
        "  labels = np.asarray([label] * num_samples)\n",
        "\n",
        "samples = sampler.get_samples(z, labels)\n",
        "imshow(imgrid(samples, cols=num_cols))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "vCffdVZvTtxL"
      },
      "outputs": [],
      "source": [
        "# @title Interpolation { run: \"auto\" }\n",
        "\n",
        "num_samples = 1  # @param {type: \"slider\", min: 1, max: 6, step: 1}\n",
        "num_interps = 6  # @param {type: \"slider\", min: 2, max: 10, step: 1}\n",
        "noise_seed_A = 11  # @param {type: \"slider\", min: 0, max: 100, step: 1}\n",
        "noise_seed_B = 0  # @param {type: \"slider\", min: 0, max: 100, step: 1}\n",
        "label_str = \"1) goldfish, Carassius auratus\"  # @param [\"0) tench, Tinca tinca\", \"1) goldfish, Carassius auratus\", \"2) great white shark, white shark, man-eater, man-eating shark, Carcharodon carcharias\", \"3) tiger shark, Galeocerdo cuvieri\", \"4) hammerhead, hammerhead shark\", \"5) electric ray, crampfish, numbfish, torpedo\", \"6) stingray\", \"7) cock\", \"8) hen\", \"9) ostrich, Struthio camelus\", \"10) brambling, Fringilla montifringilla\", \"11) goldfinch, Carduelis carduelis\", \"12) house finch, linnet, Carpodacus mexicanus\", \"13) junco, snowbird\", \"14) indigo bunting, indigo finch, indigo bird, Passerina cyanea\", \"15) robin, American robin, Turdus migratorius\", \"16) bulbul\", \"17) jay\", \"18) magpie\", \"19) chickadee\", \"20) water ouzel, dipper\", \"21) kite\", \"22) bald eagle, American eagle, Haliaeetus leucocephalus\", \"23) vulture\", \"24) great grey owl, great gray owl, Strix nebulosa\", \"25) European fire salamander, Salamandra salamandra\"]\n",
        "\n",
        "\n",
        "def interpolate(A, B, num_interps):\n",
        "  alphas = np.linspace(0, 1, num_interps)\n",
        "  if A.shape != B.shape:\n",
        "    raise ValueError('A and B must have the same shape to interpolate.')\n",
        "  return np.array([((1-a)*A + a*B)/np.sqrt(a**2 + (1-a)**2) for a in alphas])\n",
        "\n",
        "\n",
        "def interpolate_and_shape(A, B, num_interps):\n",
        "  interps = interpolate(A, B, num_interps)\n",
        "  return (interps.transpose(1, 0, *range(2, len(interps.shape)))\n",
        "                 .reshape(num_samples * num_interps, -1))\n",
        "\n",
        "label = int(label_str.split(')')[0])\n",
        "labels = np.asarray([label] * num_samples * num_interps)\n",
        "\n",
        "\n",
        "z_A = sampler.get_noise(num_samples, seed=noise_seed_A)\n",
        "z_B = sampler.get_noise(num_samples, seed=noise_seed_B)\n",
        "z = interpolate_and_shape(z_A, z_B, num_interps)\n",
        "\n",
        "samples = sampler.get_samples(z, labels)\n",
        "imshow(imgrid(samples, cols=num_interps))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "esW0Up95Ob6U"
      },
      "source": [
        "## Discriminator"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ButxPSq0OzgL"
      },
      "outputs": [],
      "source": [
        "disc = Discriminator(module_spec)\n",
        "\n",
        "batch_size = 4\n",
        "num_classes = 1000\n",
        "images = np.random.random(size=[batch_size] + disc.image_shape)\n",
        "labels = np.random.randint(0, num_classes, size=(batch_size))\n",
        "\n",
        "disc.predict(images, labels=labels)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "BhN1AplL0Hpv"
      ],
      "name": "s3gan_generation_with_tf_hub.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}