site/ko/tutorials/distribute/dtensor_ml_tutorial.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tce3stUlHN0L"
      },
      "source": [
        "##### Copyright 2019 The TensorFlow Authors.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu"
      },
      "outputs": [],
      "source": [
        "#@title 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",
        "# https://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."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MfBg1C5NB3X0"
      },
      "source": [
        "# DTensor를 사용한 분산 훈련\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r6P32iYYV27b"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td><a target=\"_blank\" href=\"https://www.tensorflow.org/tutorials/distribute/dtensor_ml_tutorial\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\">TensorFlow.org에서 보기</a></td>\n",
        "  <td>     <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/ko/tutorials/distribute/dtensor_ml_tutorial.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\">Google Colab에서 실행하기</a>\n",
        "</td>\n",
        "  <td><a target=\"_blank\" href=\"https://github.com/tensorflow/docs-l10n/blob/master/site/ko/tutorials/distribute/dtensor_ml_tutorial.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\">GitHub에서 소그 보기</a></td>\n",
        "  <td><a href=\"https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/ko/tutorials/distribute/dtensor_ml_tutorial.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\">노트북 다운론드하기</a></td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kiF4jjX4O1mF"
      },
      "source": [
        "## 개요 페이지\n",
        "\n",
        "DTensor는 여러 장치에 모델 훈련을 배포하여 효율성, 안정성 및 확장성을 개선하는 데 이용할 수 있습니다. DTensor 개념에 대한 자세한 내용은 [DTensor 프로그래밍 가이드](https://www.tensorflow.org/guide/dtensor_overview)를 참조하세요.\n",
        "\n",
        "이 튜토리얼에서는 DTensor를 사용하여 감정 분석 모델을 훈련합니다. 이 예에서는 세 가지 분산 훈련 방식을 보여줍니다.\n",
        "\n",
        "- 훈련 샘플이 장치로 샤딩(분할)되는 데이터 병렬 훈련\n",
        "- 모델 변수가 장치로 샤딩되는 모델 병렬 훈련\n",
        "- 입력 데이터의 기능이 장치에 샤딩되는 공간 병렬 훈련([공간 분할](https://cloud.google.com/blog/products/ai-machine-learning/train-ml-models-on-large-images-and-3d-volumes-with-spatial-partitioning-on-cloud-tpus)이라고도 함)\n",
        "\n",
        "이 튜토리얼의 훈련 부분은 [감정 분석에 대한 Kaggle 가이드](https://www.kaggle.com/code/anasofiauzsoy/yelp-review-sentiment-analysis-tensorflow-tfds/notebook) 노트북에서 아이디어를 얻었습니다. 전체 훈련 및 평가 워크플로(DTensor 제외)에 대해 알아보려면 해당 노트북을 참조하세요.\n",
        "\n",
        "이 튜토리얼에서는 다음 단계를 안내합니다.\n",
        "\n",
        "- 먼저 토큰화된 문장과 그 극성의 `tf.data.Dataset`을 얻기 위해 일부 데이터 정리로 시작합니다.\n",
        "\n",
        "- 다음으로, 사용자 정의 Dense 및 BatchNorm 레이어를 사용하여 MLP 모델을 빌드합니다. `tf.Module`을 사용하여 추론 변수를 추적합니다. 모델 생성자는 변수의 샤딩을 제어하기 위해 추가 `Layout` 인수를 사용합니다.\n",
        "\n",
        "- 훈련을 위해 먼저 `tf.experimental.dtensor`의 체크포인트 기능과 함께 데이터 병렬 훈련을 사용합니다. 그런 다음, 모델 병렬 훈련과 공간 병렬 훈련을 계속합니다.\n",
        "\n",
        "- 마지막 섹션에서는 TensorFlow 2.9에서 `tf.saved_model`과 `tf.experimental.dtensor` 간의 상호 작용을 간략하게 설명합니다.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YD80veeg7QtW"
      },
      "source": [
        "## 설치\n",
        "\n",
        "DTensor는 TensorFlow 2.9.0 릴리스의 일부입니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-RKXLJN-7Yyb"
      },
      "outputs": [],
      "source": [
        "!pip install --quiet --upgrade --pre tensorflow tensorflow-datasets"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tcxP4_Zu7ciQ"
      },
      "source": [
        "다음으로 `tensorflow` 및 `tensorflow.experimental.dtensor`를 가져옵니다. 그런 다음 8개의 가상 CPU를 사용하도록 TensorFlow를 구성합니다.\n",
        "\n",
        "이 예제에서는 CPU를 사용하지만 DTensor는 CPU, GPU 또는 TPU 장치에서 동일한 방식으로 작동합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dXcB26oP7dUd"
      },
      "outputs": [],
      "source": [
        "import tempfile\n",
        "import numpy as np\n",
        "import tensorflow_datasets as tfds\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        "from tensorflow.experimental import dtensor\n",
        "print('TensorFlow version:', tf.__version__)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oHtO6MJLUXlz"
      },
      "outputs": [],
      "source": [
        "def configure_virtual_cpus(ncpu):\n",
        "  phy_devices = tf.config.list_physical_devices('CPU')\n",
        "  tf.config.set_logical_device_configuration(phy_devices[0], [\n",
        "        tf.config.LogicalDeviceConfiguration(),\n",
        "    ] * ncpu)\n",
        "\n",
        "configure_virtual_cpus(8)\n",
        "DEVICES = [f'CPU:{i}' for i in range(8)]\n",
        "\n",
        "tf.config.list_logical_devices('CPU')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "omYd4jbF7j_I"
      },
      "source": [
        "## 데이터세트 다운로드\n",
        "\n",
        "감정 분석 모델을 훈련하기 위한 IMDB 리뷰 데이터세트를 다운로드합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fW4w4QlFVHhx"
      },
      "outputs": [],
      "source": [
        "train_data = tfds.load('imdb_reviews', split='train', shuffle_files=True, batch_size=64)\n",
        "train_data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ki3mpfi4aZH8"
      },
      "source": [
        "## 데이터 준비하기\n",
        "\n",
        "먼저 텍스트를 토큰화합니다. 여기에서는 원-핫 인코딩의 확장인 `tf.keras.layers.TextVectorization`의 `'tf_idf'` 모드를 사용합니다.\n",
        "\n",
        "- 속도를 위해 토큰 수를 1200개로 제한합니다.\n",
        "- `tf.Module`을 단순하게 유지하기 위해 훈련 전의 전처리 단계로 `TextVectorization`을 실행합니다.\n",
        "\n",
        "데이터 정리 섹션의 최종 결과는 토큰화된 텍스트가 `x`이고 레이블이 `y`인 `Dataset`입니다.\n",
        "\n",
        "**참고**: `TextVectorization`을 전처리 단계로 실행하는 것은 **일반적이지 않고 권장되지도 않습니다**. 그렇게 하면 훈련 데이터가 클라이언트 메모리에 맞는다고 가정하는 것이지만 항상 그런 것은 아니기 때문입니다.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zNpxjku_57Lg"
      },
      "outputs": [],
      "source": [
        "text_vectorization = tf.keras.layers.TextVectorization(output_mode='tf_idf', max_tokens=1200, output_sequence_length=None)\n",
        "text_vectorization.adapt(data=train_data.map(lambda x: x['text']))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "q16bjngoVwQp"
      },
      "outputs": [],
      "source": [
        "def vectorize(features):\n",
        "  return text_vectorization(features['text']), features['label']\n",
        "\n",
        "train_data_vec = train_data.map(vectorize)\n",
        "train_data_vec"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "atTqL9kE5wz4"
      },
      "source": [
        "## DTensor로 신경망 구축하기\n",
        "\n",
        "이제 `DTensor`를 사용하여 MLP(Multi-Layer Perceptron) 네트워크를 구축합니다. 이 네트워크는 완전히 연결된 Dense 및 BatchNorm 레이어를 사용합니다.\n",
        "\n",
        "`DTensor`는 입력 `Tensor` 및 변수의 `dtensor.Layout` 속성에 따라 일반 TensorFlow Ops의 SPMD(단일 프로그램 다중 데이터) 확장을 통해 TensorFlow를 확장합니다.\n",
        "\n",
        "`DTensor` 인식 레이어의 변수는 `dtensor.DVariable`이며 `DTensor` 인식 레이어 객체의 생성자는 일반적인 레이어 매개변수 외에 추가 `Layout` 입력을 받습니다.\n",
        "\n",
        "참고: TensorFlow 2.9부터 tf.keras.layer.Dense 및 `tf.keras.layer.BatchNormalization`과 같은 `tf.keras.layer.Dense` 레이어는 `dtensor.Layout` 인수를 허용합니다. Keras를 DTensor와 함께 사용하기 위한 자세한 내용은 [DTensor Keras 통합 튜토리얼](/tutorials/distribute/dtensor_keras_tutorial)을 참조하세요."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PMCt-Gj3b3Jy"
      },
      "source": [
        "### Dense 레이어\n",
        "\n",
        "다음 사용자 정의 Dense 레이어는 2개의 레이어 변수를 정의합니다. $W_{ij}$는 가중치에 대한 변수이고 $b_i$는 바이어스에 대한 변수입니다.\n",
        "\n",
        "$$\n",
        "y_j = \\sigma(\\sum_i x_i W_{ij} + b_j)\n",
        "$$\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nYlFUJWNjl4N"
      },
      "source": [
        "### 레이아웃 추론\n",
        "\n",
        "이 결과는 다음 관찰로부터 비롯됩니다.\n",
        "\n",
        "- 행렬 내적 $t_j = \\sum_i x_i W_{ij}$에 대한 피연산자의 기본 DTensor 샤딩은 $i$ 축을 따라 $\\mathbf{W}$ 및 $\\mathbf{x}$를 같은 방식으로 샤딩하는 것입니다.\n",
        "\n",
        "- 행렬 합 $t_j + b_j$에 대한 피연산자에 대한 기본 DTensor 샤딩은 $j$ 축을 따라 동일한 방식으로 $\\mathbf{t}$ 및 $\\mathbf{b}$를 분할하는 것입니다.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VpKblz7Yb16G"
      },
      "outputs": [],
      "source": [
        "class Dense(tf.Module):\n",
        "\n",
        "  def __init__(self, input_size, output_size,\n",
        "               init_seed, weight_layout, activation=None):\n",
        "    super().__init__()\n",
        "\n",
        "    random_normal_initializer = tf.function(tf.random.stateless_normal)\n",
        "\n",
        "    self.weight = dtensor.DVariable(\n",
        "        dtensor.call_with_layout(\n",
        "            random_normal_initializer, weight_layout,\n",
        "            shape=[input_size, output_size],\n",
        "            seed=init_seed\n",
        "            ))\n",
        "    if activation is None:\n",
        "      activation = lambda x:x\n",
        "    self.activation = activation\n",
        "    \n",
        "    # bias is sharded the same way as the last axis of weight.\n",
        "    bias_layout = weight_layout.delete([0])\n",
        "\n",
        "    self.bias = dtensor.DVariable(\n",
        "        dtensor.call_with_layout(tf.zeros, bias_layout, [output_size]))\n",
        "\n",
        "  def __call__(self, x):\n",
        "    y = tf.matmul(x, self.weight) + self.bias\n",
        "    y = self.activation(y)\n",
        "\n",
        "    return y"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tfVY_vAKbxM0"
      },
      "source": [
        "### BatchNorm\n",
        "\n",
        "배치 정규화 레이어는 훈련 중 모드 축소를 방지하는 데 도움이 됩니다. 이 경우 배치 정규화 레이어를 추가하면 모델 훈련이 0만 생성하는 모델을 생성하는 것을 방지할 수 있습니다.\n",
        "\n",
        "아래 사용자 정의 `BatchNorm` 레이어의 생성자는 `Layout` 인수를 사용하지 않습니다. `BatchNorm`에는 레이어 변수가 없기 때문입니다. 그래도 DTensor에서의 작동에는 문제가 없는데 레이어에 대한 유일한 입력인 'x'가 이미 전역 배치를 나타내는 DTensor이기 때문입니다.\n",
        "\n",
        "참고: DTensor에서 입력 Tensor 'x'는 항상 전역 배치를 나타냅니다. 따라서 `tf.nn.batch_normalization`은 전역 배치에 적용됩니다. `tf.distribute.MirroredStrategy`를 사용한 훈련은 이와 다르며, 여기서는 Tensor 'x'가 배치(로컬 배치)의 복제본당 샤드만 나타냅니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "riBA9pfhlPFq"
      },
      "outputs": [],
      "source": [
        "class BatchNorm(tf.Module):\n",
        "\n",
        "  def __init__(self):\n",
        "    super().__init__()\n",
        "\n",
        "  def __call__(self, x, training=True):\n",
        "    if not training:\n",
        "      # This branch is not used in the Tutorial.\n",
        "      pass\n",
        "    mean, variance = tf.nn.moments(x, axes=[0])\n",
        "    return tf.nn.batch_normalization(x, mean, variance, 0.0, 1.0, 1e-5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q4R4MPz5prh4"
      },
      "source": [
        "완전한 기능을 갖춘 배치 정규화 레이어(예: `tf.keras.layers.BatchNormalization`)에는 변수에 대한 레이아웃 인수가 필요합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "unFcP99zprJj"
      },
      "outputs": [],
      "source": [
        "def make_keras_bn(bn_layout):\n",
        "  return tf.keras.layers.BatchNormalization(gamma_layout=bn_layout,\n",
        "                                            beta_layout=bn_layout,\n",
        "                                            moving_mean_layout=bn_layout,\n",
        "                                            moving_variance_layout=bn_layout,\n",
        "                                            fused=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v8Dj7AJ_lPs0"
      },
      "source": [
        "### 레이어 합치기\n",
        "\n",
        "다음으로, 위의 구성 요소를 사용하여 MLP(Multi-layer perceptron) 네트워크를 구축합니다. 아래 다이어그램은 DTensor 샤딩 또는 복제가 적용되지 않은 두 `Dense` 레이어에 대한 입력 `x`와 가중치 행렬 간의 축 관계를 보여줍니다."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "udFGAO-NrZw6"
      },
      "source": [
        "<img src=\"https://www.tensorflow.org/images/dtensor/no_dtensor.png\" class=\"no-filter\" alt=\"The input and weight matrices for a non distributed model.\">\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8DCQ0aQ5rQtB"
      },
      "source": [
        "첫 번째 `Dense` 레이어의 출력은 `BatchNorm` 이후의 두 번째 `Dense` 레이어의 입력으로 전달됩니다. 따라서 첫 번째 `Dense` 레이어($\\mathbf{W_1}$)의 출력과 두 번째 Dense 레이어($\\mathbf{W_2}$)의 입력에 대한 기본 DTensor 샤딩은 $\\hat{j}$ 공통 축을 따라 같은 방식으로 $\\mathbf{W_1}$ 및 $\\mathbf{W_2}$를 샤딩하는 것입니다.\n",
        "\n",
        "$$\n",
        "\\mathsf{Layout}[{W_{1,ij}}; i, j] = \\left[\\hat{i}, \\hat{j}\\right] \\\\\n",
        "\\mathsf{Layout}[{W_{2,jk}}; j, k] = \\left[\\hat{j}, \\hat{k} \\right]\n",
        "$$\n",
        "\n",
        "레이아웃 추론이 2개의 레이아웃이 독립적이지 않음을 보여주더라도 모델 인터페이스를 단순화하기 위해 `MLP`는 Dense 레이어당 하나씩 2개의 `Layout` 인수를 사용합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "junyS-965opl"
      },
      "outputs": [],
      "source": [
        "from typing import Tuple\n",
        "\n",
        "class MLP(tf.Module):\n",
        "\n",
        "  def __init__(self, dense_layouts: Tuple[dtensor.Layout, dtensor.Layout]):\n",
        "    super().__init__()\n",
        "\n",
        "    self.dense1 = Dense(\n",
        "        1200, 48, (1, 2), dense_layouts[0], activation=tf.nn.relu)\n",
        "    self.bn = BatchNorm()\n",
        "    self.dense2 = Dense(48, 2, (3, 4), dense_layouts[1])\n",
        "\n",
        "  def __call__(self, x):\n",
        "    y = x\n",
        "    y = self.dense1(y)\n",
        "    y = self.bn(y)\n",
        "    y = self.dense2(y)\n",
        "    return y\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9dgLmebHhr7h"
      },
      "source": [
        "레이아웃 추론 제약의 정확성과 API의 단순성 사이의 절충은 DTensor를 사용하는 API 설계 시 공통적인 주안점입니다. 다른 API를 사용하여 `Layout` 간의 종속성을 포착하는 것도 가능합니다. 예를 들어 `MLPStricter` 클래스는 생성자에서 `Layout` 객체를 만듭니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wEZR7UlihsYX"
      },
      "outputs": [],
      "source": [
        "class MLPStricter(tf.Module):\n",
        "\n",
        "  def __init__(self, mesh, input_mesh_dim, inner_mesh_dim1, output_mesh_dim):\n",
        "    super().__init__()\n",
        "\n",
        "    self.dense1 = Dense(\n",
        "        1200, 48, (1, 2), dtensor.Layout([input_mesh_dim, inner_mesh_dim1], mesh),\n",
        "        activation=tf.nn.relu)\n",
        "    self.bn = BatchNorm()\n",
        "    self.dense2 = Dense(48, 2, (3, 4), dtensor.Layout([inner_mesh_dim1, output_mesh_dim], mesh))\n",
        "\n",
        "\n",
        "  def __call__(self, x):\n",
        "    y = x\n",
        "    y = self.dense1(y)\n",
        "    y = self.bn(y)\n",
        "    y = self.dense2(y)\n",
        "    return y"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GcQi7D5mal2L"
      },
      "source": [
        "모델이 실행되는지 확인하려면 완전히 복제된 레이아웃과 `'x'` 입력의 완전히 복제된 배치로 모델을 검사해야 합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zOPuYeQwallh"
      },
      "outputs": [],
      "source": [
        "WORLD = dtensor.create_mesh([(\"world\", 8)], devices=DEVICES)\n",
        "\n",
        "model = MLP([dtensor.Layout.replicated(WORLD, rank=2),\n",
        "             dtensor.Layout.replicated(WORLD, rank=2)])\n",
        "\n",
        "sample_x, sample_y = train_data_vec.take(1).get_single_element()\n",
        "sample_x = dtensor.copy_to_mesh(sample_x, dtensor.Layout.replicated(WORLD, rank=2))\n",
        "print(model(sample_x))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "akrjDstEpDv9"
      },
      "source": [
        "## 장치로 데이터 이동하기\n",
        "\n",
        "일반적으로 `tf.data` 반복기(및 기타 데이터 가져오기 방법)는 로컬 호스트 장치 메모리가 지원하는 텐서 객체를 생성합니다. 이 데이터는 DTensor의 구성 요소 텐서를 지원하는 가속기 장치 메모리로 전송되어야 합니다.\n",
        "\n",
        "`dtensor.copy_to_mesh`는 이 상황에 적합하지 않은데, DTensor의 전역 관점으로 인해 입력 텐서를 모든 장치에 복제하기 때문입니다. 따라서 이 튜토리얼에서는 원활한 데이터 전송을 위해 헬퍼 함수 `repack_local_tensor`를 사용합니다. 이 헬퍼 함수는 `dtensor.pack`을 사용하여 복제본을 지원하는 장치에 대한 복제본의 용도로서 전역 배치의 샤드를 보냅니다(보내기만 함).\n",
        "\n",
        "이 단순화된 기능은 단일 클라이언트를 가정합니다. 로컬 텐서를 분할하는 올바른 방법과 분할 조각과 로컬 장치 간의 매핑을 결정하는 것은 다중 클라이언트 애플리케이션에서 힘들 수 있습니다.\n",
        "\n",
        "`tf.data` 통합을 단순화하여 단일 클라이언트와 다중 클라이언트 애플리케이션을 모두 지원하기 위한 추가 DTensor API가 계획되어 있습니다. 계속 지켜봐 주세요."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3t5WvQR4Hvo4"
      },
      "outputs": [],
      "source": [
        "def repack_local_tensor(x, layout):\n",
        "  \"\"\"Repacks a local Tensor-like to a DTensor with layout.\n",
        "\n",
        "  This function assumes a single-client application.\n",
        "  \"\"\"\n",
        "  x = tf.convert_to_tensor(x)\n",
        "  sharded_dims = []\n",
        "\n",
        "  # For every sharded dimension, use tf.split to split the along the dimension.\n",
        "  # The result is a nested list of split-tensors in queue[0].\n",
        "  queue = [x]\n",
        "  for axis, dim in enumerate(layout.sharding_specs):\n",
        "    if dim == dtensor.UNSHARDED:\n",
        "      continue\n",
        "    num_splits = layout.shape[axis]\n",
        "    queue = tf.nest.map_structure(lambda x: tf.split(x, num_splits, axis=axis), queue)\n",
        "    sharded_dims.append(dim)\n",
        "\n",
        "  # Now we can build the list of component tensors by looking up the location in\n",
        "  # the nested list of split-tensors created in queue[0].\n",
        "  components = []\n",
        "  for locations in layout.mesh.local_device_locations():\n",
        "    t = queue[0]\n",
        "    for dim in sharded_dims:\n",
        "      split_index = locations[dim]  # Only valid on single-client mesh.\n",
        "      t = t[split_index]\n",
        "    components.append(t)\n",
        "\n",
        "  return dtensor.pack(components, layout)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2KKCDcjG7zj2"
      },
      "source": [
        "## 데이터 병렬 훈련\n",
        "\n",
        "이 섹션에서는 데이터 병렬 훈련으로 MLP 모델을 훈련합니다. 다음 섹션에서는 모델 병렬 훈련과 공간 병렬 훈련을 보여줍니다.\n",
        "\n",
        "데이터 병렬 훈련은 분산 머린 러닝에 일반적으로 사용되는 방식입니다.\n",
        "\n",
        "- 모델 변수는 각각 N개의 장치에 복제됩니다.\n",
        "- 전역 배치는 복제본당 N개의 배치로 분할됩니다.\n",
        "- 각 복제본의 배치는 복제본 장치에서 훈련됩니다.\n",
        "- 모든 복제본에서 데이터 가중치 부여가 일괄적으로 수행되기 전에 그래디언트가 감소합니다.\n",
        "\n",
        "데이터 병렬 훈련은 장치 수와 관련하여 거의 선형적인 속도 향상을 제공합니다."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UMsLUyTGq3oL"
      },
      "source": [
        "### 데이터 병렬 메쉬 생성하기\n",
        "\n",
        "일반적인 데이터 병렬 처리 훈련 루프는 단일 `batch` 차원으로 구성된 DTensor `Mesh`를 사용하며, 여기서 각 장치는 전역 배치에서 샤드를 받는 복제본이 됩니다.\n",
        "\n",
        "<img src=\"https://www.tensorflow.org/images/dtensor/dtensor_data_para.png\" class=\"no-filter\" alt=\"Data parallel mesh\">\n",
        "\n",
        "복제된 모델은 복제본에서 실행되므로 모델 변수는 완전히 복제됩니다(샤딩되지 않음)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "C0IyOlxmeu4I"
      },
      "outputs": [],
      "source": [
        "mesh = dtensor.create_mesh([(\"batch\", 8)], devices=DEVICES)\n",
        "\n",
        "model = MLP([dtensor.Layout([dtensor.UNSHARDED, dtensor.UNSHARDED], mesh),\n",
        "             dtensor.Layout([dtensor.UNSHARDED, dtensor.UNSHARDED], mesh),])\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OREKwBybo1gZ"
      },
      "source": [
        "### DTensor에 훈련 데이터 패킹하기\n",
        "\n",
        "훈련 데이터 배치는 `'batch'`(첫 번째) 축을 따라 샤딩된 DTensor로 패킹되어야 하며, 그래야 DTensor가 훈련 데이터를 `'batch'` 메쉬 차원에 고르게 분배합니다.\n",
        "\n",
        "**참고**: DTensor에서 `batch size`는 항상 전역 배치 크기를 나타냅니다. 배치 크기는 `batch` 메쉬 차원의 크기로 균등하게 나눠질 수 있도록 선택해야 합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8xMYkTpGocY8"
      },
      "outputs": [],
      "source": [
        "def repack_batch(x, y, mesh):\n",
        "  x = repack_local_tensor(x, layout=dtensor.Layout(['batch', dtensor.UNSHARDED], mesh))\n",
        "  y = repack_local_tensor(y, layout=dtensor.Layout(['batch'], mesh))\n",
        "  return x, y\n",
        "\n",
        "sample_x, sample_y = train_data_vec.take(1).get_single_element()\n",
        "sample_x, sample_y = repack_batch(sample_x, sample_y, mesh)\n",
        "\n",
        "print('x', sample_x[:, 0])\n",
        "print('y', sample_y)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uONSiqOIkFL1"
      },
      "source": [
        "### 훈련 스텝\n",
        "\n",
        "이 예에서는 사용자 정의 훈련 루프(CTL)와 함께 확률적 경사 하강법 옵티마이저를 사용합니다. 이러한 주제에 대한 자세한 내용은 [사용자 정의 훈련 루프 가이드](https://www.tensorflow.org/guide/keras/writing_a_training_loop_from_scratch) 및 [둘러보기](https://www.tensorflow.org/tutorials/customization/custom_training_walkthrough)를 참조하세요.\n",
        "\n",
        "`train_step`은 `tf.function`으로 캡슐화되어 이 본문이 TensorFlow 그래프로 추적될 것임을 나타냅니다. `train_step`의 본문은 순방향 추론 패스, 역방향 그래디언트 패스 및 변수 업데이트로 구성됩니다.\n",
        "\n",
        "`train_step`의 본문에는 특별한 DTensor 주석이 포함되어 있지 않습니다. 대신 `train_step`에는 입력 배치 및 모델의 전역 보기에서 입력 `x` 및 `y`를 처리하는 상위 수준 TensorFlow 연산만 포함됩니다. 모든 DTensor 주석( `Mesh`, `Layout`)은 훈련 스텝에서 고려되지 않습니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BwUFzLGDtQT6"
      },
      "outputs": [],
      "source": [
        "# Refer to the CTL (custom training loop guide)\n",
        "@tf.function\n",
        "def train_step(model, x, y, learning_rate=tf.constant(1e-4)):\n",
        "  with tf.GradientTape() as tape:\n",
        "    logits = model(x)\n",
        "    # tf.reduce_sum sums the batch sharded per-example loss to a replicated\n",
        "    # global loss (scalar).\n",
        "    loss = tf.reduce_sum(\n",
        "        tf.nn.sparse_softmax_cross_entropy_with_logits(\n",
        "            logits=logits, labels=y))\n",
        "  parameters = model.trainable_variables\n",
        "  gradients = tape.gradient(loss, parameters)\n",
        "  for parameter, parameter_gradient in zip(parameters, gradients):\n",
        "    parameter.assign_sub(learning_rate * parameter_gradient)\n",
        "\n",
        "  # Define some metrics\n",
        "  accuracy = 1.0 - tf.reduce_sum(tf.cast(tf.argmax(logits, axis=-1, output_type=tf.int64) != y, tf.float32)) / x.shape[0]\n",
        "  loss_per_sample = loss / len(x)\n",
        "  return {'loss': loss_per_sample, 'accuracy': accuracy}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0OYTu4j0evWT"
      },
      "source": [
        "### 체크포인트\n",
        "\n",
        "`tf.train.Checkpoint`를 사용하여 DTensor 모델을 즉시 체크포인트할 수 있습니다. 샤딩된 DVariables를 저장하고 복원하면 효율적인 샤딩된 저장과 복원이 수행됩니다. 현재, `tf.train.Checkpoint.save` 및 `tf.train.Checkpoint.restore`를 사용할 때 모든 DVariables는 동일한 호스트 메시에 있어야 하며 DVariables와 일반 변수는 함께 저장할 수 없습니다. [이 가이드](../../guide/checkpoint.ipynb)에서 체크포인트에 대해 자세히 알아볼 수 있습니다.\n",
        "\n",
        "DTensor 체크포인트가 복원되면 변수의 `Layout`이 체크포인트가 저장될 때와 다를 수 있습니다. 즉, DTensor 모델의 저장은 레이아웃 및 메시에 구애받지 않으며 샤딩된 저장의 효율성에만 영향을 미칩니다. 하나의 메시와 레이아웃으로 DTensor 모델을 저장하고 이를 다른 메시와 레이아웃에 복원할 수 있습니다. 이 튜토리얼에서는 이 기능을 사용하여 모델 병렬 훈련과 공간 병렬 훈련 섹션에서 훈련을 계속합니다.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rsInFFJg7x9t"
      },
      "outputs": [],
      "source": [
        "CHECKPOINT_DIR = tempfile.mkdtemp()\n",
        "\n",
        "def start_checkpoint_manager(model):\n",
        "  ckpt = tf.train.Checkpoint(root=model)\n",
        "  manager = tf.train.CheckpointManager(ckpt, CHECKPOINT_DIR, max_to_keep=3)\n",
        "\n",
        "  if manager.latest_checkpoint:\n",
        "    print(\"Restoring a checkpoint\")\n",
        "    ckpt.restore(manager.latest_checkpoint).assert_consumed()\n",
        "  else:\n",
        "    print(\"New training\")\n",
        "  return manager\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9r77ky5Jgp1j"
      },
      "source": [
        "### 훈련 루프\n",
        "\n",
        "데이터 병렬 훈련 방식의 경우, epoch에 대해 훈련하고 진행 상황을 보고합니다. 3개의 epoch로는 모델 훈련에 충분하지 않습니다. 50%의 정확도는 무작위로 추측하는 정도에 불과합니다.\n",
        "\n",
        "나중에 훈련을 할 수 있도록 체크포인트를 활성화합니다. 다음 섹션에서는 체크포인트를 로드하고 다른 병렬 방식으로 훈련합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UaLn-vGZgqbS"
      },
      "outputs": [],
      "source": [
        "num_epochs = 2\n",
        "manager = start_checkpoint_manager(model)\n",
        "\n",
        "for epoch in range(num_epochs):\n",
        "  step = 0\n",
        "  pbar = tf.keras.utils.Progbar(target=int(train_data_vec.cardinality()), stateful_metrics=[])\n",
        "  metrics = {'epoch': epoch}\n",
        "  for x,y in train_data_vec:\n",
        "\n",
        "    x, y = repack_batch(x, y, mesh)\n",
        "\n",
        "    metrics.update(train_step(model, x, y, 1e-2))\n",
        "\n",
        "    pbar.update(step, values=metrics.items(), finalize=False)\n",
        "    step += 1\n",
        "  manager.save()\n",
        "  pbar.update(step, values=metrics.items(), finalize=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YRFJEhum7EGD"
      },
      "source": [
        "## 모델 병렬 훈련\n",
        "\n",
        "2차원 `Mesh`로 전환하고 두 번째 메쉬 차원을 따라 모델 변수를 샤딩하면 훈련이 모델 병렬이 됩니다.\n",
        "\n",
        "모델 병렬 훈련에서 각 모델 복제본은 여러 장치(이 경우 2개)에 걸쳐 있습니다.\n",
        "\n",
        "- 4개의 모델 복제본이 있으며 훈련 데이터 배치는 4개의 복제본에 배포됩니다.\n",
        "- 단일 모델 복제본 내의 2개 장치는 복제된 훈련 데이터를 수신합니다.\n",
        "\n",
        "<img src=\"https://www.tensorflow.org/images/dtensor/dtensor_model_para.png\" class=\"no-filter\" alt=\"Model parallel mesh\">\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5gZE9IT5Dzwl"
      },
      "outputs": [],
      "source": [
        "mesh = dtensor.create_mesh([(\"batch\", 4), (\"model\", 2)], devices=DEVICES)\n",
        "model = MLP([dtensor.Layout([dtensor.UNSHARDED, \"model\"], mesh), \n",
        "             dtensor.Layout([\"model\", dtensor.UNSHARDED], mesh)])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ihof3DkMFKnf"
      },
      "source": [
        "훈련 데이터는 여전히 배치 차원을 따라 샤딩되므로 데이터 병렬 훈련 사례와 동일한 `repack_batch` 함수를 재사용할 수 있습니다. DTensor는 `\"model\"` 메쉬 차원을 따라 복제본 내부의 모든 장치에 복제본당 배치를 자동으로 복제합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dZf56ynbE_p1"
      },
      "outputs": [],
      "source": [
        "def repack_batch(x, y, mesh):\n",
        "  x = repack_local_tensor(x, layout=dtensor.Layout(['batch', dtensor.UNSHARDED], mesh))\n",
        "  y = repack_local_tensor(y, layout=dtensor.Layout(['batch'], mesh))\n",
        "  return x, y"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UW3OXdhNFfpv"
      },
      "source": [
        "다음으로, 훈련 루프를 실행합니다. 훈련 루프는 데이터 병렬 훈련 예제와 동일한 체크포인트 관리자를 재사용하며 코드는 동일해 보입니다.\n",
        "\n",
        "모델 병렬 학습에서 데이터 병렬 학습 모델을 계속 훈련할 수 있습니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LLC0wgii7EgA"
      },
      "outputs": [],
      "source": [
        "num_epochs = 2\n",
        "manager = start_checkpoint_manager(model)\n",
        "\n",
        "for epoch in range(num_epochs):\n",
        "  step = 0\n",
        "  pbar = tf.keras.utils.Progbar(target=int(train_data_vec.cardinality()))\n",
        "  metrics = {'epoch': epoch}\n",
        "  for x,y in train_data_vec:\n",
        "    x, y = repack_batch(x, y, mesh)\n",
        "    metrics.update(train_step(model, x, y, 1e-2))\n",
        "    pbar.update(step, values=metrics.items(), finalize=False)\n",
        "    step += 1\n",
        "  manager.save()\n",
        "  pbar.update(step, values=metrics.items(), finalize=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BZH-aMrVzi2L"
      },
      "source": [
        "## 공간 병렬 훈련"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u-bK6IZ9GCS9"
      },
      "source": [
        "매우 높은 차원의 데이터(예: 매우 큰 이미지 또는 비디오)를 훈련할 때 특성 차원을 따라 샤딩하는 것이 바람직할 수 있습니다. 이를 [공간 분할(Spatial Partitioning)](https://cloud.google.com/blog/products/ai-machine-learning/train-ml-models-on-large-images-and-3d-volumes-with-spatial-partitioning-on-cloud-tpus)이라고 하며, 이는 큰 3차원 입력 샘플이 있는 모델을 훈련하기 위해 TensorFlow에 처음 도입되었습니다.\n",
        "\n",
        "<img src=\"https://www.tensorflow.org/images/dtensor/dtensor_spatial_para.png\" class=\"no-filter\" alt=\"Spatial parallel mesh\">\n",
        "\n",
        "DTensor도 이러한 사례를 지원합니다. 유일하게 변경해야 하는 부분은 `feature` 차원을 포함하는 메쉬를 생성하고 해당 `Layout`을 적용하는 것입니다.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jpc9mqURGpmK"
      },
      "outputs": [],
      "source": [
        "mesh = dtensor.create_mesh([(\"batch\", 2), (\"feature\", 2), (\"model\", 2)], devices=DEVICES)\n",
        "model = MLP([dtensor.Layout([\"feature\", \"model\"], mesh), \n",
        "             dtensor.Layout([\"model\", dtensor.UNSHARDED], mesh)])\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i07Wrv-jHBc1"
      },
      "source": [
        "입력 텐서를 DTensor로 패킹할 때 `feature` 차원을 따라 입력 데이터를 샤딩합니다. 약간 다른 repack 함수인 `repack_batch_for_spt`를 사용하여 이 작업을 수행합니다. 여기서 `spt`는 Spatial Parallel Training을 나타냅니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DWR8qF6BGtFL"
      },
      "outputs": [],
      "source": [
        "def repack_batch_for_spt(x, y, mesh):\n",
        "    # Shard data on feature dimension, too\n",
        "    x = repack_local_tensor(x, layout=dtensor.Layout([\"batch\", 'feature'], mesh))\n",
        "    y = repack_local_tensor(y, layout=dtensor.Layout([\"batch\"], mesh))\n",
        "    return x, y"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ygl9dqMUHTVN"
      },
      "source": [
        "다른 병렬 훈련 방식으로 생성된 체크포인트에서 공간 병렬 훈련을 계속 이어갈 수도 있습니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "p3NnpHSKo-hx"
      },
      "outputs": [],
      "source": [
        "num_epochs = 2\n",
        "\n",
        "manager = start_checkpoint_manager(model)\n",
        "for epoch in range(num_epochs):\n",
        "  step = 0\n",
        "  metrics = {'epoch': epoch}\n",
        "  pbar = tf.keras.utils.Progbar(target=int(train_data_vec.cardinality()))\n",
        "\n",
        "  for x, y in train_data_vec:\n",
        "    x, y = repack_batch_for_spt(x, y, mesh)\n",
        "    metrics.update(train_step(model, x, y, 1e-2))\n",
        "\n",
        "    pbar.update(step, values=metrics.items(), finalize=False)\n",
        "    step += 1\n",
        "  manager.save()\n",
        "  pbar.update(step, values=metrics.items(), finalize=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vp4L59CpJjYr"
      },
      "source": [
        "## SavedModel 및 DTensor\n",
        "\n",
        "DTensor와 SavedModel의 통합은 아직 개발 중입니다.\n",
        "\n",
        "TensorFlow `2.11`을 기준으로 `tf.saved_model`은 샤딩 및 복제된 DTensor 모델을 저장할 수 있으며 저장은 메쉬의 서로 다른 장치에서 효율적인 샤딩 저장을 수행합니다. 그러나 모델이 저장되면 모든 DTensor 주석이 손실되고 저장된 서명은 DTensor가 아닌 일반 Tensor에서만 사용할 수 있습니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "49HfIq_SJZoj"
      },
      "outputs": [],
      "source": [
        "mesh = dtensor.create_mesh([(\"world\", 1)], devices=DEVICES[:1])\n",
        "mlp = MLP([dtensor.Layout([dtensor.UNSHARDED, dtensor.UNSHARDED], mesh), \n",
        "           dtensor.Layout([dtensor.UNSHARDED, dtensor.UNSHARDED], mesh)])\n",
        "\n",
        "manager = start_checkpoint_manager(mlp)\n",
        "\n",
        "model_for_saving = tf.keras.Sequential([\n",
        "  text_vectorization,\n",
        "  mlp\n",
        "])\n",
        "\n",
        "@tf.function(input_signature=[tf.TensorSpec([None], tf.string)])\n",
        "def run(inputs):\n",
        "  return {'result': model_for_saving(inputs)}\n",
        "\n",
        "tf.saved_model.save(\n",
        "    model_for_saving, \"/tmp/saved_model\",\n",
        "    signatures=run)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h6Csim_VMGxQ"
      },
      "source": [
        "TensorFlow 2.9.0부터 일반 Tensor 또는 완전히 복제된 DTensor(일반 Tensor로 변환됨)로만 로드된 서명을 호출할 수 있습니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HG_ASSzR4IWW"
      },
      "outputs": [],
      "source": [
        "sample_batch = train_data.take(1).get_single_element()\n",
        "sample_batch"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qW8yKPrhKQ5b"
      },
      "outputs": [],
      "source": [
        "loaded = tf.saved_model.load(\"/tmp/saved_model\")\n",
        "\n",
        "run_sig = loaded.signatures[\"serving_default\"]\n",
        "result = run_sig(sample_batch['text'])['result']"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GahGbv0ZmkJb"
      },
      "outputs": [],
      "source": [
        "np.mean(tf.argmax(result, axis=-1) == sample_batch['label'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ks-Vs9qsH6jO"
      },
      "source": [
        "## 다음은?\n",
        "\n",
        "이 튜토리얼에서는 DTensor를 사용하여 MLP 감정 분석 모델을 빌드하고 훈련하는 방법을 보여주었습니다.\n",
        "\n",
        "`Mesh` 및 `Layout` 기본 형식을 통해 DTensor는 TensorFlow `tf.function`을 다양한 훈련 방식에 적합한 분산 프로그램으로 변환할 수 있습니다.\n",
        "\n",
        "실제 머신 러닝 애플리케이션에서는 과적합된 모델을 생성하지 않도록 평가와 교차 검증을 적용해야 합니다. 이 튜토리얼에서 소개된 방법은 평가에 병렬 처리를 도입하는 데도 적용할 수 있습니다.\n",
        "\n",
        "처음부터 `tf.Module`로 모델을 구성하려면 많은 작업이 필요하지만 레이어 및 헬퍼 함수와 같은 기존 빌딩 블록을 재사용하면 모델 개발 속도를 크게 높일 수 있습니다. TensorFlow 2.9부터 `tf.keras.layers` 아래의 모든 Keras 레이어는 DTensor 레이아웃을 인수로 허용하며 DTensor 모델을 빌드하는 데 사용할 수 있습니다. 모델 구현을 수정하지 않고 DTensor로 Keras 모델을 직접 재사용할 수도 있습니다. DTensor Keras 사용에 대한 정보는 [DTensor Keras 통합 튜토리얼](https://www.tensorflow.org/tutorials/distribute/dtensor_keras_tutorial)을 참조하세요. "
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "dtensor_ml_tutorial.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}