ch02_7_JAX에서의_병렬_평가.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "V28",
      "authorship_tag": "ABX9TyPVxyll7gqAere/opTN3tSY",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "TPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/yblee110/jax-flax-book/blob/main/ch02_7_JAX%EC%97%90%EC%84%9C%EC%9D%98_%EB%B3%91%EB%A0%AC_%ED%8F%89%EA%B0%80.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install jax==0.4.26\n",
        "import jax\n",
        "print(jax.__version__)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "k0YaiRvnkVhW",
        "outputId": "59936633-d0c2-4bee-d4a4-7ea039226954"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: jax==0.4.26 in /usr/local/lib/python3.10/dist-packages (0.4.26)\n",
            "Requirement already satisfied: ml-dtypes>=0.2.0 in /usr/local/lib/python3.10/dist-packages (from jax==0.4.26) (0.2.0)\n",
            "Requirement already satisfied: numpy>=1.22 in /usr/local/lib/python3.10/dist-packages (from jax==0.4.26) (1.25.2)\n",
            "Requirement already satisfied: opt-einsum in /usr/local/lib/python3.10/dist-packages (from jax==0.4.26) (3.3.0)\n",
            "Requirement already satisfied: scipy>=1.9 in /usr/local/lib/python3.10/dist-packages (from jax==0.4.26) (1.11.4)\n",
            "0.4.26\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "9hP_9qf4rGlX"
      },
      "outputs": [],
      "source": [
        "import jax.tools.colab_tpu\n",
        "jax.tools.colab_tpu.setup_tpu()"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import jax\n",
        "import jax.numpy as jnp\n",
        "jax.devices()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AsWpSQpjrNsu",
        "outputId": "ad4f6ba4-b827-4846-f182-70c43718f179"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0),\n",
              " TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1),\n",
              " TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0),\n",
              " TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1),\n",
              " TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0),\n",
              " TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1),\n",
              " TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0),\n",
              " TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1)]"
            ]
          },
          "metadata": {},
          "execution_count": 3
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "\n",
        "x = np.arange(5)\n",
        "w = np.array([2., 3., 4.])\n",
        "\n",
        "def convolve(x, w):\n",
        "  output = []\n",
        "  for i in range(1, len(x)-1):\n",
        "    output.append(jnp.dot(x[i-1:i+2], w))\n",
        "  return jnp.array(output)\n",
        "\n",
        "convolve(x, w)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XR0Nd1QGroMH",
        "outputId": "0a63b4b8-0082-429a-ba00-17c371a548f1"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([11., 20., 29.], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 4
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "n_devices = jax.local_device_count()\n",
        "xs = np.arange(5 * n_devices).reshape(-1, 5)\n",
        "ws = np.stack([w] * n_devices)\n",
        "\n",
        "xs"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vLAJNTBdrujS",
        "outputId": "40afd553-d459-499f-c308-5358245a4bc9"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([[ 0,  1,  2,  3,  4],\n",
              "       [ 5,  6,  7,  8,  9],\n",
              "       [10, 11, 12, 13, 14],\n",
              "       [15, 16, 17, 18, 19],\n",
              "       [20, 21, 22, 23, 24],\n",
              "       [25, 26, 27, 28, 29],\n",
              "       [30, 31, 32, 33, 34],\n",
              "       [35, 36, 37, 38, 39]])"
            ]
          },
          "metadata": {},
          "execution_count": 5
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "ws"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wlHf_gXxryV6",
        "outputId": "88711f3a-80d9-48c6-ac46-b2ebd5cf9864"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "array([[2., 3., 4.],\n",
              "       [2., 3., 4.],\n",
              "       [2., 3., 4.],\n",
              "       [2., 3., 4.],\n",
              "       [2., 3., 4.],\n",
              "       [2., 3., 4.],\n",
              "       [2., 3., 4.],\n",
              "       [2., 3., 4.]])"
            ]
          },
          "metadata": {},
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jax.vmap(convolve)(xs, ws)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "x6HRgmy4r4ZO",
        "outputId": "39576f86-ee14-4659-cea8-92bdaf244835"
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([[ 11.,  20.,  29.],\n",
              "       [ 56.,  65.,  74.],\n",
              "       [101., 110., 119.],\n",
              "       [146., 155., 164.],\n",
              "       [191., 200., 209.],\n",
              "       [236., 245., 254.],\n",
              "       [281., 290., 299.],\n",
              "       [326., 335., 344.]], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "jax.pmap(convolve)(xs, ws)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Jnwwb-F6r6Oc",
        "outputId": "ae87fdb7-e89e-4b91-b3bc-d4c9043cf5ac"
      },
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Array([[ 11.,  20.,  29.],\n",
              "       [ 56.,  65.,  74.],\n",
              "       [101., 110., 119.],\n",
              "       [146., 155., 164.],\n",
              "       [191., 200., 209.],\n",
              "       [236., 245., 254.],\n",
              "       [281., 290., 299.],\n",
              "       [326., 335., 344.]], dtype=float32)"
            ]
          },
          "metadata": {},
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import time\n",
        "start = time.time()\n",
        "jax.vmap(convolve)(xs, ws)\n",
        "print(f\"vmap : {time.time()-start:.4f} sec\")\n",
        "start = time.time()\n",
        "jax.pmap(convolve)(xs, ws)\n",
        "print(f\"pmap : {time.time()-start:.4f} sec\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cGZhJYdZr73n",
        "outputId": "9d88cf09-0277-4df4-95cf-ffba8ad1ed80"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "vmap : 0.0074 sec\n",
            "pmap : 0.0021 sec\n"
          ]
        }
      ]
    }
  ]
}