site/ko/guide/core/distribution.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FhGuhbZ6M5tl"
      },
      "source": [
        "##### Copyright 2022 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "AwOEIRJC6Une"
      },
      "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": "EIdT9iu_Z4Rb"
      },
      "source": [
        "# Core API 및 DTensor를 사용하는 분산 훈련"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bBIlTPscrIT9"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>     <a target=\"_blank\" href=\"https://www.tensorflow.org/guide/core/distribution\"><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/guide/core/distribution.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\">Google Colab에서 실행하기</a> </td>\n",
        "  <td><a target=\"_blank\" href=\"https://github.com/tensorflow/docs-l10n/blob/master/site/ko/guide/core/distribution.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/guide/core/distribution.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\">노트북 다운로드하기</a></td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SjAxxRpBzVYg"
      },
      "source": [
        "## 소개\n",
        "\n",
        "이 노트북은 [TensorFlow Core 하위 수준 API](https://www.tensorflow.org/guide/core) 및 [DTensor](https://www.tensorflow.org/guide/dtensor_overview)를 사용하는 데이터 병렬 분산 훈련 예제를 보여줍니다. TensorFlow Core 및 기본 사용 사례에 대해 자세히 알아보려면 [Core API 개요](https://www.tensorflow.org/guide/core)를 확인하세요. DTensor에 대한 자세한 내용은 [DTensor 개요](https://www.tensorflow.org/guide/dtensor_overview) 가이드와 [DTensor를 사용하는 분산 훈련](https://www.tensorflow.org/tutorials/distribute/dtensor_ml_tutorial) 가이드를 참고하세요.\n",
        "\n",
        "이 예제에서는 [멀티 레이어 퍼셉트론](https://www.tensorflow.org/guide/core/mlp_core) 가이드와 동일한 모델 및 최적화 도구를 사용합니다. Core API를 사용하는 엔드 투 엔드 머신러닝 워크플로를 작성하는 데 익숙해지려면 먼저 이 튜토리얼을 참고하세요.\n",
        "\n",
        "참고: DTensor는 아직 실험적인 TensorFlow API입니다. 즉, 테스트에서 해당 기능을 사용할 수 있습니다. 사실 테스트 환경에서만 사용할 수 있습니다."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d_OFkG0dyWCp"
      },
      "source": [
        "## DTensor를 사용하는 데이터 병렬 훈련 개요\n",
        "\n",
        "분산 기능을 지원하는 MLP를 빌드하기 전에 잠시 시간을 내어 데이터 병렬 훈련용 DTensor의 기본 사항을 살펴보도록 합니다.\n",
        "\n",
        "DTensor를 사용하여 장치 전반에서 분산 훈련을 실행하고 효율성과 안정성, 그리고 확장성을 개선할 수 있습니다. DTensor는 SPMD(Single program, multiple data) 확장이라는 절차와 샤딩 지시문에 따라 프로그램과 텐서를 배포합니다. `DTensor` 인식 레이어의 변수는 `dtensor.DVariable`로 생성되고, `DTensor` 인식 레이어 객체의 생성자는 일반적인 레이어 매개변수에 `Layout` 입력을 추가적으로 더합니다.\n",
        "\n",
        "데이터 병렬 훈련의 주요 개념은 다음과 같습니다.\n",
        "\n",
        "- 모델 변수는 각각 N개의 장치에 복제됩니다.\n",
        "- 전역 배치는 복제본당 N개의 배치로 분할됩니다.\n",
        "- 각 복제본의 배치는 복제본 장치에서 훈련됩니다.\n",
        "- 모든 복제본에서 데이터 가중치 부여가 일괄적으로 수행되기 전에 그래디언트가 감소합니다.\n",
        "- 데이터 병렬 훈련은 장치 수와 거의 비례하는 선형 속도를 제공합니다."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nchsZfwEVtVs"
      },
      "source": [
        "## 설치하기\n",
        "\n",
        "DTensor는 TensorFlow 2.9.0 릴리스에 포함된 항목입니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "latuqlI_Yvoo"
      },
      "outputs": [],
      "source": [
        "#!pip install --quiet --upgrade --pre tensorflow"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1rRo8oNqZ-Rj"
      },
      "outputs": [],
      "source": [
        "import matplotlib\n",
        "from matplotlib import pyplot as plt\n",
        "# Preset Matplotlib figure sizes.\n",
        "matplotlib.rcParams['figure.figsize'] = [9, 6]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9xQKvCJ85kCQ"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "from tensorflow.experimental import dtensor\n",
        "print(tf.__version__)\n",
        "# Set random seed for reproducible results \n",
        "tf.random.set_seed(22)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vDH9-sy4sfPf"
      },
      "source": [
        "이 실험을 위해 8개의 가상 CPU를 구성합니다. DTensor는 GPU 또는 TPU 장치와 함께 사용할 수도 있습니다. 이 노트북에서는 가상 장치를 사용하기 때문에 분산 훈련으로 인한 속도 향상은 눈에 띄지 않습니다. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "H2iM-6J4s2D6"
      },
      "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",
        "\n",
        "DEVICES = [f'CPU:{i}' for i in range(8)]\n",
        "devices = tf.config.list_logical_devices('CPU')\n",
        "device_names = [d.name for d in devices]\n",
        "device_names"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F_72b0LCNbjx"
      },
      "source": [
        "## MNIST 데이터세트\n",
        "\n",
        "[TensorFlow 데이터세트](https://www.tensorflow.org/datasets/catalog/mnist)의 데이터세트를 사용할 수 있습니다. 데이터는 훈련 세트와 테스트 세트로 나눕니다. 시간을 절약하기 위해 훈련과 테스트에 5000개의 예제만 사용합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8h4fV_JCfPIX"
      },
      "outputs": [],
      "source": [
        "train_data, test_data = tfds.load(\"mnist\", split=['train[:5000]', 'test[:5000]'], batch_size=128, as_supervised=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "twkJ35YB6tSi"
      },
      "source": [
        "### 데이터 전처리하기\n",
        "\n",
        "데이터를 2차원으로 변형하고 단위 간격 [0,1]에 맞도록 크기를 다시 조정하여 데이터를 전처리합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6Cmjhg0xCqbz"
      },
      "outputs": [],
      "source": [
        "def preprocess(x, y):\n",
        "  # Reshaping the data\n",
        "  x = tf.reshape(x, shape=[-1, 784])\n",
        "  # Rescaling the data\n",
        "  x = x/255\n",
        "  return x, y\n",
        "\n",
        "train_data, test_data = train_data.map(preprocess), test_data.map(preprocess)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6o3CrycBXA2s"
      },
      "source": [
        "## MLP 빌드하기\n",
        "\n",
        "DTensor 인식 레이어로 MLP 모델을 빌드합니다."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OHW6Yvg2yS6H"
      },
      "source": [
        "### 밀집 레이어\n",
        "\n",
        "DTensor를 지원하는 밀집 레이어 모듈을 생성하는 것으로 시작합니다. `dtensor.call_with_layout` 함수는 DTensor 입력을 받고 DTensor 출력을 생성하는 함수를 호출하는 데 사용할 수 있습니다. 이것은 TensorFlow 지원 함수로 DTensor 변수 `dtensor.DVariable`을 초기화하는 데 유용합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IM0yJos25FG5"
      },
      "outputs": [],
      "source": [
        "class DenseLayer(tf.Module):\n",
        "\n",
        "  def __init__(self, in_dim, out_dim, weight_layout, activation=tf.identity):\n",
        "    super().__init__()\n",
        "    # Initialize dimensions and the activation function\n",
        "    self.in_dim, self.out_dim = in_dim, out_dim\n",
        "    self.activation = activation\n",
        "\n",
        "    # Initialize the DTensor weights using the Xavier scheme\n",
        "    uniform_initializer = tf.function(tf.random.stateless_uniform)\n",
        "    xavier_lim = tf.sqrt(6.)/tf.sqrt(tf.cast(self.in_dim + self.out_dim, tf.float32))\n",
        "    self.w = dtensor.DVariable(\n",
        "      dtensor.call_with_layout(\n",
        "          uniform_initializer, weight_layout,\n",
        "          shape=(self.in_dim, self.out_dim), seed=(22, 23),\n",
        "          minval=-xavier_lim, maxval=xavier_lim))\n",
        "        \n",
        "    # Initialize the bias with the zeros\n",
        "    bias_layout = weight_layout.delete([0])\n",
        "    self.b = dtensor.DVariable(\n",
        "      dtensor.call_with_layout(tf.zeros, bias_layout, shape=[out_dim]))\n",
        "\n",
        "  def __call__(self, x):\n",
        "    # Compute the forward pass\n",
        "    z = tf.add(tf.matmul(x, self.w), self.b)\n",
        "    return self.activation(z)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X-7MzpjgyHg6"
      },
      "source": [
        "### MLP 순차 모델\n",
        "\n",
        "이제 밀집 레이어를 순차적으로 실행하는 MLP 모듈을 생성합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6XisRWiCyHAb"
      },
      "outputs": [],
      "source": [
        "class MLP(tf.Module):\n",
        "\n",
        "  def __init__(self, layers):\n",
        "    self.layers = layers\n",
        "   \n",
        "  def __call__(self, x, preds=False): \n",
        "    # Execute the model's layers sequentially\n",
        "    for layer in self.layers:\n",
        "      x = layer(x)\n",
        "    return x"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r5HZJ0kv-V3v"
      },
      "source": [
        "DTensor로 \"데이터 병렬\" 훈련을 수행하는 작업은 `tf.distribute.MirroredStrategy`와 동일합니다. 이를 위해 각 장치는 데이터 배치의 샤드에서 동일한 모델을 실행합니다. 따라서 다음이 필요합니다.\n",
        "\n",
        "- 단일 `\"batch\"` 차원의 `dtensor.Mesh`\n",
        "- 메시 전체에 가중치를 복제하는 모든 가중치용 `dtensor.Layout`(각 축에 `dtensor.UNSHARDED` 사용)\n",
        "- 메시에서 배치 차원을 분할하는 데이터용 `dtensor.Layout`\n",
        "\n",
        "단일 배치 차원으로 구성된 DTensor 메시를 생성합니다. 여기서 각 장치는 전역 배치에서 샤드를 수신하는 복제본이 됩니다. 이 메시를 사용하여 다음 아키텍처로 MLP 모드를 인스턴스화합니다.\n",
        "\n",
        "순방향 전달: ReLU(784 x 700) x ReLU(700 x 500) x Softmax(500 x 10)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VmlACuki3oPi"
      },
      "outputs": [],
      "source": [
        "mesh = dtensor.create_mesh([(\"batch\", 8)], devices=DEVICES)\n",
        "weight_layout = dtensor.Layout([dtensor.UNSHARDED, dtensor.UNSHARDED], mesh)\n",
        "\n",
        "input_size = 784\n",
        "hidden_layer_1_size = 700\n",
        "hidden_layer_2_size = 500\n",
        "hidden_layer_2_size = 10\n",
        "\n",
        "mlp_model = MLP([\n",
        "    DenseLayer(in_dim=input_size, out_dim=hidden_layer_1_size, \n",
        "               weight_layout=weight_layout,\n",
        "               activation=tf.nn.relu),\n",
        "    DenseLayer(in_dim=hidden_layer_1_size , out_dim=hidden_layer_2_size,\n",
        "               weight_layout=weight_layout,\n",
        "               activation=tf.nn.relu),\n",
        "    DenseLayer(in_dim=hidden_layer_2_size, out_dim=hidden_layer_2_size, \n",
        "               weight_layout=weight_layout)])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tyBATDoRmDkg"
      },
      "source": [
        "### 훈련 메트릭\n",
        "\n",
        "훈련에 교차 엔트로피 손실 함수와 정확성 메트릭을 사용합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rskOYA7FVCwg"
      },
      "outputs": [],
      "source": [
        "def cross_entropy_loss(y_pred, y):\n",
        "  # Compute cross entropy loss with a sparse operation\n",
        "  sparse_ce = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=y, logits=y_pred)\n",
        "  return tf.reduce_mean(sparse_ce)\n",
        "\n",
        "def accuracy(y_pred, y):\n",
        "  # Compute accuracy after extracting class predictions\n",
        "  class_preds = tf.argmax(y_pred, axis=1)\n",
        "  is_equal = tf.equal(y, class_preds)\n",
        "  return tf.reduce_mean(tf.cast(is_equal, tf.float32))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JSiNRhTOnKZr"
      },
      "source": [
        "### 옵티마이저\n",
        "\n",
        "옵티마이저를 사용하면 표준 경사 하강법에 비해 훨씬 더 빠른 수렴 결과를 얻을 수 있습니다. Adam 옵티마이저가 아래에 구현되어 있으며DTensor와 호환되도록 구성되어 있습니다. Keras 옵티마이저를 DTensor와 함께 사용하려면 실험적 `tf.keras.dtensor.experimental.optimizers` 모듈을 참고합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-9kIAI_lfXDS"
      },
      "outputs": [],
      "source": [
        "class Adam(tf.Module):\n",
        "\n",
        "    def __init__(self, model_vars, learning_rate=1e-3, beta_1=0.9, beta_2=0.999, ep=1e-7):\n",
        "      # Initialize optimizer parameters and variable slots\n",
        "      self.model_vars = model_vars\n",
        "      self.beta_1 = beta_1\n",
        "      self.beta_2 = beta_2\n",
        "      self.learning_rate = learning_rate\n",
        "      self.ep = ep\n",
        "      self.t = 1.\n",
        "      self.v_dvar, self.s_dvar = [], []\n",
        "      # Initialize optimizer variable slots\n",
        "      for var in model_vars:\n",
        "        v = dtensor.DVariable(dtensor.call_with_layout(tf.zeros, var.layout, shape=var.shape))\n",
        "        s = dtensor.DVariable(dtensor.call_with_layout(tf.zeros, var.layout, shape=var.shape))\n",
        "        self.v_dvar.append(v)\n",
        "        self.s_dvar.append(s)\n",
        "\n",
        "    def apply_gradients(self, grads):\n",
        "      # Update the model variables given their gradients\n",
        "      for i, (d_var, var) in enumerate(zip(grads, self.model_vars)):\n",
        "        self.v_dvar[i].assign(self.beta_1*self.v_dvar[i] + (1-self.beta_1)*d_var)\n",
        "        self.s_dvar[i].assign(self.beta_2*self.s_dvar[i] + (1-self.beta_2)*tf.square(d_var))\n",
        "        v_dvar_bc = self.v_dvar[i]/(1-(self.beta_1**self.t))\n",
        "        s_dvar_bc = self.s_dvar[i]/(1-(self.beta_2**self.t))\n",
        "        var.assign_sub(self.learning_rate*(v_dvar_bc/(tf.sqrt(s_dvar_bc) + self.ep)))\n",
        "      self.t += 1.\n",
        "      return "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w54b7GtLfn1j"
      },
      "source": [
        "### 데이터 패킹\n",
        "\n",
        "장치에 데이터를 전송하는 헬퍼 함수를 작성하는 것으로 시작합니다. 이 함수는 `dtensor.pack`을 사용하여 복제본을 지원하는 장치로 복제본을 대상으로 하는 전역 배치의 샤드를 보내야(전송만 해야 함) 합니다. 편리한 진행을 위해 단일 클라이언트 애플리케이션을 가정합니다.\n",
        "\n",
        "다음으로 이 헬퍼 함수를 사용하여 훈련 데이터 배치를 배치축(첫 번째)을 따라 샤딩한 DTensor로 패킹하는 함수를 작성합니다. 이렇게 하면 DTensor가 훈련 데이터를 '배치' 메시 차원에 고르게 분배합니다. DTensor에서 배치 크기는 항상 전역 배치 크기를 나타냅니다. 따라서 배치 크기는 배치 메시 차원의 크기로 균등하게 나눌 수 있도록 선택해야 합니다. `tf.data` 통합을 단순화하기 위한 추가 DTensor API가 계획되어 있으므로 계속 지켜봐 주시기 바랍니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3Rx82djZ6ITm"
      },
      "outputs": [],
      "source": [
        "def repack_local_tensor(x, layout):\n",
        "  # Repacks a local Tensor-like to a DTensor with layout\n",
        "  # This function assumes a single-client application\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 you 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)\n",
        "\n",
        "def repack_batch(x, y, mesh):\n",
        "  # Pack training data batches into DTensors along the batch axis\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": "osEK3rqpYfKd"
      },
      "source": [
        "### 훈련하기\n",
        "\n",
        "데이터 배치를 제공할 경우 단일 훈련 단계를 실행하는 추적 가능한 함수를 작성합니다. 이 함수는 특별한 DTensor 주석을 요구하지 않습니다. 또한 테스트 단계를 실행하고 적절한 성능 메트릭을 반환하는 함수를 작성하도록 합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZICEsDGuSbDD"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def train_step(model, x_batch, y_batch, loss, metric, optimizer):\n",
        "  # Execute a single training step\n",
        "  with tf.GradientTape() as tape:\n",
        "    y_pred = model(x_batch)\n",
        "    batch_loss = loss(y_pred, y_batch)\n",
        "  # Compute gradients and update the model's parameters\n",
        "  grads = tape.gradient(batch_loss, model.trainable_variables)\n",
        "  optimizer.apply_gradients(grads)\n",
        "  # Return batch loss and accuracy\n",
        "  batch_acc = metric(y_pred, y_batch)\n",
        "  return batch_loss, batch_acc\n",
        "\n",
        "@tf.function\n",
        "def test_step(model, x_batch, y_batch, loss, metric):\n",
        "  # Execute a single testing step\n",
        "  y_pred = model(x_batch)\n",
        "  batch_loss = loss(y_pred, y_batch)\n",
        "  batch_acc = metric(y_pred, y_batch)\n",
        "  return batch_loss, batch_acc"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RjIDVTwwX-Mr"
      },
      "source": [
        "이제 배치 크기가 128인 3 epoch의 MLP 모델을 훈련합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oC85kuZgmh3q"
      },
      "outputs": [],
      "source": [
        "# Initialize the training loop parameters and structures\n",
        "epochs = 3\n",
        "batch_size = 128\n",
        "train_losses, test_losses = [], []\n",
        "train_accs, test_accs = [], []\n",
        "optimizer = Adam(mlp_model.trainable_variables)\n",
        "\n",
        "# Format training loop\n",
        "for epoch in range(epochs):\n",
        "  batch_losses_train, batch_accs_train = [], []\n",
        "  batch_losses_test, batch_accs_test = [], []\n",
        "\n",
        "  # Iterate through training data\n",
        "  for x_batch, y_batch in train_data:\n",
        "    x_batch, y_batch = repack_batch(x_batch, y_batch, mesh)\n",
        "    batch_loss, batch_acc = train_step(mlp_model, x_batch, y_batch, cross_entropy_loss, accuracy, optimizer)\n",
        "   # Keep track of batch-level training performance\n",
        "    batch_losses_train.append(batch_loss)\n",
        "    batch_accs_train.append(batch_acc)\n",
        "\n",
        "  # Iterate through testing data\n",
        "  for x_batch, y_batch in test_data:\n",
        "    x_batch, y_batch = repack_batch(x_batch, y_batch, mesh)\n",
        "    batch_loss, batch_acc = test_step(mlp_model, x_batch, y_batch, cross_entropy_loss, accuracy)\n",
        "    # Keep track of batch-level testing\n",
        "    batch_losses_test.append(batch_loss)\n",
        "    batch_accs_test.append(batch_acc)\n",
        "\n",
        "# Keep track of epoch-level model performance\n",
        "  train_loss, train_acc = tf.reduce_mean(batch_losses_train), tf.reduce_mean(batch_accs_train)\n",
        "  test_loss, test_acc = tf.reduce_mean(batch_losses_test), tf.reduce_mean(batch_accs_test)\n",
        "  train_losses.append(train_loss)\n",
        "  train_accs.append(train_acc)\n",
        "  test_losses.append(test_loss)\n",
        "  test_accs.append(test_acc)\n",
        "  print(f\"Epoch: {epoch}\")\n",
        "  print(f\"Training loss: {train_loss.numpy():.3f}, Training accuracy: {train_acc.numpy():.3f}\")\n",
        "  print(f\"Testing loss: {test_loss.numpy():.3f}, Testing accuracy: {test_acc.numpy():.3f}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j_RVmt43G12R"
      },
      "source": [
        "### 성능 평가\n",
        "\n",
        "먼저 훈련하는 동안 모델의 손실과 정확성을 시각화하는 플로팅 함수를 작성합니다. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VXTCYVtNDjAM"
      },
      "outputs": [],
      "source": [
        "def plot_metrics(train_metric, test_metric, metric_type):\n",
        "  # Visualize metrics vs training Epochs\n",
        "  plt.figure()\n",
        "  plt.plot(range(len(train_metric)), train_metric, label = f\"Training {metric_type}\")\n",
        "  plt.plot(range(len(test_metric)), test_metric, label = f\"Testing {metric_type}\")\n",
        "  plt.xlabel(\"Epochs\")\n",
        "  plt.ylabel(metric_type)\n",
        "  plt.legend()\n",
        "  plt.title(f\"{metric_type} vs Training Epochs\");"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "407qok7q2JIO"
      },
      "outputs": [],
      "source": [
        "plot_metrics(train_losses, test_losses, \"Cross entropy loss\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8H_TgxV92NfX"
      },
      "outputs": [],
      "source": [
        "plot_metrics(train_accs, test_accs, \"Accuracy\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DHO_u-3w4YRF"
      },
      "source": [
        "## 모델 저장하기\n",
        "\n",
        "`tf.saved_model`과 DTensor의 통합은 아직 개발을 진행하고 있습니다. TensorFlow 2.9.0부터 tf.saved_model은 완전히 복제된 변수가 있는 DTensor 모델만 허용합니다. 해결 방법으로 체크포인트를 다시 로드하여 DTensor 모델을 완전히 복제된 모델로 전환할 수 있습니다. 그러나 모델을 저장한 후에는 모든 DTensor 주석이 손실되며 저장된 서명은 일반 Tensor로만 사용할 수 있습니다. 안정화되면 통합을 보여주기 위해 이 튜토리얼이 업데이트됩니다.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VFLfEH4ManbW"
      },
      "source": [
        "## 결론\n",
        "\n",
        "이 노트북에서는 DTensor 및 TensorFlow Core API를 사용하는 분산 훈련의 개요를 제공했습니다. 다음은 도움이 될 수 있는 몇 가지 추가 정보입니다.\n",
        "\n",
        "- 분산 훈련을 지원하고 고도로 구성 가능한 머신러닝 워크플로를 구축하는 데 [TensorFlow Core API](https://www.tensorflow.org/guide/core)를 사용할 수 있습니다.\n",
        "- [DTensor 개념](https://www.tensorflow.org/guide/dtensor_overview) 가이드 및 [DTensor를 사용하는 분산 훈련](https://www.tensorflow.org/tutorials/distribute/dtensor_ml_tutorial) 가이드에는 DTensor 및 통합에 대한 최신 정보가 포함되어 있습니다.\n",
        "\n",
        "TensorFlow Core API를 사용하는 더 많은 예제는 [가이드](https://www.tensorflow.org/guide/core)를 확인하세요. 데이터 로드 및 준비에 대해 자세히 알아보려면 [이미지 데이터 로드](https://www.tensorflow.org/tutorials/load_data/images) 또는 [CSV 데이터 로드](https://www.tensorflow.org/tutorials/load_data/csv) 튜토리얼을 참고하세요."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "FhGuhbZ6M5tl"
      ],
      "name": "distribution.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}