kfp-dist-train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b784c29d-e0b6-4d42-adc7-c3ae39c57ad6",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q kfp==1.7.1 kfpdist  -i https://pypi.tuna.tsinghua.edu.cn/simple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2cf0df8d-46e7-44f6-9f56-4f63961276b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import kfp\n",
    "from kfp import dsl as v1dsl\n",
    "from kfp.v2 import dsl\n",
    "from kfp.v2.dsl import Input, InputPath, Output, OutputPath, Dataset, Model, Metrics, component\n",
    "import kfp.compiler as compiler\n",
    "import os\n",
    "from typing import List\n",
    "from kfpdist import set_dist_train_config\n",
    "\n",
    "@component\n",
    "def prepare(n_workers: int,\n",
    "            output_list_path: OutputPath(List[int])):\n",
    "    import json\n",
    "    with open(output_list_path, 'w') as fn:\n",
    "        fn.write(json.dumps(list(range(n_workers))))\n",
    "\n",
    "@component(base_image='tensorflow/tensorflow:2.5.1-gpu', packages_to_install=['tensorflow_datasets'])\n",
    "def train(rank:   int,\n",
    "          nranks: int,\n",
    "          output_model_path: OutputPath('output_model_path'),\n",
    "          output_tblog_path: OutputPath('output_tblog_path'),\n",
    "          output_ckpt_path:  OutputPath('output_ckpt_path')):\n",
    "    import tensorflow as tf\n",
    "    from tensorflow.keras import layers, models\n",
    "    import os\n",
    "    \n",
    "    set_dist_train_config(rank, nranks, 'train', port=9888)\n",
    "    \n",
    "    def make_datasets_unbatched():\n",
    "        import tensorflow_datasets as tfds\n",
    "        BUFFER_SIZE = 10000\n",
    "        # Scaling MNIST data from (0, 255] to (0., 1.]\n",
    "        def scale(image, label):\n",
    "            image = tf.cast(image, tf.float32)\n",
    "            image /= 255\n",
    "            return image, label\n",
    "        datasets, _ = tfds.load(name='mnist', with_info=True, as_supervised=True)\n",
    "\n",
    "        return datasets['train'].map(scale).cache().shuffle(BUFFER_SIZE)\n",
    "    \n",
    "    def build_and_compile_cnn_model():\n",
    "        model = models.Sequential()\n",
    "        model.add(\n",
    "          layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)))\n",
    "        model.add(layers.MaxPooling2D((2, 2)))\n",
    "        model.add(layers.Conv2D(64, (3, 3), activation='relu'))\n",
    "        model.add(layers.MaxPooling2D((2, 2)))\n",
    "        model.add(layers.Conv2D(64, (3, 3), activation='relu'))\n",
    "        model.add(layers.Flatten())\n",
    "        model.add(layers.Dense(64, activation='relu'))\n",
    "        model.add(layers.Dense(10, activation='softmax'))\n",
    "        model.summary()\n",
    "        model.compile(optimizer='adam',\n",
    "                      loss='sparse_categorical_crossentropy',\n",
    "                      metrics=['accuracy'])\n",
    "        return model\n",
    "\n",
    "    def decay(epoch):\n",
    "        if epoch < 3:\n",
    "            return 1e-3\n",
    "        if 3 <= epoch < 7:\n",
    "            return 1e-4\n",
    "        return 1e-5\n",
    "\n",
    "    # Use NCCL if you need to use GPU\n",
    "    strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy(\n",
    "      communication=tf.distribute.experimental.CollectiveCommunication.RING)\n",
    "    BATCH_SIZE_PER_REPLICA = 64\n",
    "    BATCH_SIZE = BATCH_SIZE_PER_REPLICA * strategy.num_replicas_in_sync\n",
    "\n",
    "    with strategy.scope():\n",
    "        ds_train = make_datasets_unbatched_s3().batch(BATCH_SIZE).repeat()\n",
    "        options = tf.data.Options()\n",
    "        options.experimental_distribute.auto_shard_policy = \\\n",
    "            tf.data.experimental.AutoShardPolicy.DATA\n",
    "        ds_train = ds_train.with_options(options)\n",
    "        multi_worker_model = build_and_compile_cnn_model()\n",
    "\n",
    "    checkpoint_prefix = os.path.join(output_ckpt_path, \"ckpt_{epoch}\")\n",
    "\n",
    "    class PrintLR(tf.keras.callbacks.Callback):\n",
    "        def on_epoch_end(self, epoch, logs=None): #pylint: disable=no-self-use\n",
    "            print('\\nLearning rate for epoch {} is {}'.format(\n",
    "                epoch + 1, multi_worker_model.optimizer.lr.numpy()))\n",
    "\n",
    "    callbacks = [\n",
    "        tf.keras.callbacks.TensorBoard(log_dir=output_tblog_path),\n",
    "        tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_prefix,\n",
    "                                           save_weights_only=True),\n",
    "        tf.keras.callbacks.LearningRateScheduler(decay),\n",
    "        PrintLR()\n",
    "    ]\n",
    "    multi_worker_model.fit(ds_train,\n",
    "                           epochs=10,\n",
    "                           steps_per_epoch=70,\n",
    "                           callbacks=callbacks)\n",
    "\n",
    "    def is_chief():\n",
    "        return rank == 0\n",
    "\n",
    "    if is_chief():\n",
    "        model_path = output_model_path + '/v1'\n",
    "    else:\n",
    "        model_path = output_model_path + '/worker_tmp_' + str(rank)\n",
    "    multi_worker_model.save(model_path)\n",
    "\n",
    "@dsl.pipeline(pipeline_root='', name='quickstart-pipeline')\n",
    "def dist_train_pipeline(n_workers:       int = 2,\n",
    "                        use_gpu:         int = 0,\n",
    "                        gpus_per_worker: int = 1,\n",
    "                        num_epochs:      int = 15,\n",
    "                        train_batchsize: int = 128,\n",
    "                        test_batchsize:  int = 128,\n",
    "                        learning_rate:   float = 0.001,\n",
    "                        model_version:   int = 1):\n",
    "    prep_step = prepare(n_workers)\n",
    "    with dsl.ParallelFor(prep_step.outputs['output_list']) as rank:\n",
    "        train_step = train(rank, n_workers)\n",
    "        # Uncomment below line if you need to use GPU.\n",
    "        # train_step.set_gpu_limit(1)\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    limg = \"typhoon1986/ml-pipeline-kfp-launcher:1.7.1\"\n",
    "    kfp.Client().create_run_from_pipeline_func(dist_train_pipeline,\n",
    "                                               launcher_image=limg,\n",
    "                                               mode=v1dsl.PipelineExecutionMode.V2_COMPATIBLE,\n",
    "                                               arguments={})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7460b146-aa6c-41f8-a828-f5960312508d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}