04-Model Training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "33b0bbba",
   "metadata": {},
   "source": [
    "# 04 - Model Training\n",
    "\n",
    "This notebook creates a custom model training container component that is deployed to Artifact Registry."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "24f3587b-1321-43dd-a4e9-24db68815810",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "import logging\n",
    "import kfp\n",
    "from google.cloud import bigquery, storage\n",
    "from google.cloud import aiplatform as vertex_ai\n",
    "from google_cloud_pipeline_components.experimental.custom_job import utils\n",
    "from kfp.v2 import compiler, dsl\n",
    "from kfp.v2.dsl import component\n",
    "from typing import NamedTuple\n",
    "from kfp.v2.dsl import (Artifact, Dataset, Input, InputPath, Model, Output, Metrics,\n",
    "                        OutputPath, component)\n",
    "\n",
    "from google_cloud_pipeline_components.experimental.custom_job import utils"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be9f73ad",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "993bb57e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PIPELINE_ROOT: gs://mle_airbus_dataset/airbusmlepipeline/pipeline_root\n",
      "MODULE_ROOT: gs://mle_airbus_dataset/airbusmlepipeline/pipeline_module\n",
      "DATA_ROOT: gs://mle_airbus_dataset/airbusmlepipeline/data\n",
      "SERVING_MODEL_DIR: gs://mle_airbus_dataset/airbusmlepipeline/serving_model\n",
      "Project ID: mle-airbus-detection-smu\n",
      "Region: asia-east1\n",
      "Bucket name: mle_airbus_dataset\n",
      "Service Account: service-account-for-mle@mle-airbus-detection-smu.iam.gserviceaccount.com\n",
      "Vertex API Parent URI: projects/mle-airbus-detection-smu/locations/asia-east1\n",
      "      2373  2022-06-26T04:02:36Z  gs://mle_airbus_dataset/mle-airbus-detection-smu-b1f8ee58e814.json#1656216156596465  metageneration=1\n",
      "    194297  2022-07-03T10:12:15Z  gs://mle_airbus_dataset/test.parquet#1656843135799925  metageneration=1\n",
      "   1541531  2022-07-03T10:12:15Z  gs://mle_airbus_dataset/train.parquet#1656843135725932  metageneration=1\n",
      " 419430400  2022-07-03T10:13:34Z  gs://mle_airbus_dataset/train_hist.csv#1656843214414169  metageneration=1\n",
      "        18  2022-07-03T10:13:34Z  gs://mle_airbus_dataset/train_threshold.txt#1656843214485788  metageneration=1\n",
      "                                 gs://mle_airbus_dataset/airbusmlepipeline/\n",
      "                                 gs://mle_airbus_dataset/pipeline_root/\n",
      "                                 gs://mle_airbus_dataset/train_v2/\n",
      "                                 gs://mle_airbus_dataset/trained_model/\n",
      "TOTAL: 5 objects, 421168619 bytes (401.66 MiB)\n",
      "gs://mle_airbus_dataset/train_v2/\n"
     ]
    }
   ],
   "source": [
    "PROJECT = 'mle-airbus-detection-smu' # Change to your project id.\n",
    "PROJECT_NUMBER = '484894607141'\n",
    "REGION = 'asia-east1' # Change to your region.\n",
    "BUCKET = 'mle_airbus_dataset' # Change to your bucket name.\n",
    "SERVICE_ACCOUNT = \"service-account-for-mle@mle-airbus-detection-smu.iam.gserviceaccount.com\"\n",
    "\n",
    "if PROJECT == \"\" or PROJECT is None or PROJECT == \"[your-project-id]\":\n",
    "    # Get your GCP project id from gcloud\n",
    "    shell_output = !gcloud config list --format 'value(core.project)' 2>/dev/null\n",
    "    PROJECT = shell_output[0]\n",
    "    \n",
    "if SERVICE_ACCOUNT == \"\" or SERVICE_ACCOUNT is None or SERVICE_ACCOUNT == \"[your-service-account]\":\n",
    "    # Get your GCP project id from gcloud\n",
    "    shell_output = !gcloud config list --format 'value(core.account)' 2>/dev/null\n",
    "    SERVICE_ACCOUNT = shell_output[0]\n",
    "    \n",
    "if BUCKET == \"\" or BUCKET is None or BUCKET == \"[your-bucket-name]\":\n",
    "    # Get your bucket name to GCP projet id\n",
    "    BUCKET = PROJECT\n",
    "    # Try to create the bucket if it doesn'exists\n",
    "    ! gsutil mb -l $REGION gs://$BUCKET\n",
    "    print(\"\")\n",
    "    \n",
    "PARENT = f\"projects/{PROJECT}/locations/{REGION}\"\n",
    "\n",
    "PIPELINE_NAME = 'airbusmlepipeline'\n",
    "\n",
    "# Path to various pipeline artifact.\n",
    "PIPELINE_ROOT = 'gs://{}/{}/pipeline_root'.format(\n",
    "    BUCKET, PIPELINE_NAME)\n",
    "\n",
    "# Paths for users' Python module.\n",
    "MODULE_ROOT = 'gs://{}/{}/pipeline_module'.format(\n",
    "    BUCKET, PIPELINE_NAME)\n",
    "\n",
    "# Paths for users' data.\n",
    "DATA_ROOT = 'gs://{}/{}/data'.format(BUCKET, PIPELINE_NAME)\n",
    "\n",
    "# This is the path where your model will be pushed for serving.\n",
    "SERVING_MODEL_DIR = 'gs://{}/{}/serving_model'.format(\n",
    "    BUCKET, PIPELINE_NAME)\n",
    "\n",
    "print('PIPELINE_ROOT: {}'.format(PIPELINE_ROOT))\n",
    "print('MODULE_ROOT: {}'.format(MODULE_ROOT))\n",
    "print('DATA_ROOT: {}'.format(DATA_ROOT))\n",
    "print('SERVING_MODEL_DIR: {}'.format(SERVING_MODEL_DIR))\n",
    "\n",
    "BQ_DATASET_NAME = 'mle-airbus-detection-smu.airbus_label_dataset' # Change to your BQ dataset name.\n",
    "BQ_TABLE_NAME = 'airbus_label'\n",
    "BQ_LOCATION = ' asia-east1'\n",
    "    \n",
    "print(\"Project ID:\", PROJECT)\n",
    "print(\"Region:\", REGION)\n",
    "print(\"Bucket name:\", BUCKET)\n",
    "print(\"Service Account:\", SERVICE_ACCOUNT)\n",
    "print(\"Vertex API Parent URI:\", PARENT)\n",
    "\n",
    "! gsutil ls -al \"gs://\"$BUCKET\n",
    "\n",
    "\n",
    "storage_client = storage.Client(project=PROJECT)\n",
    "bucket = storage_client.get_bucket(BUCKET)\n",
    "storage_path = f\"gs://{BUCKET}/train_v2/\"\n",
    "print (storage_path)\n",
    "\n",
    "VERSION = 'v01'\n",
    "DATASET_DISPLAY_NAME = 'airbus-ship-dataset-display'\n",
    "MODEL_DISPLAY_NAME = f'{DATASET_DISPLAY_NAME}-classifier-{VERSION}'\n",
    "\n",
    "WORKSPACE = f'gs://{BUCKET}/{DATASET_DISPLAY_NAME}'\n",
    "EXPERIMENT_ARTIFACTS_DIR = os.path.join(WORKSPACE, 'experiments')\n",
    "\n",
    "TENSORBOARD_DISPLAY_NAME = f'tb-{DATASET_DISPLAY_NAME}'\n",
    "EXPERIMENT_NAME = f'{MODEL_DISPLAY_NAME}'\n",
    "\n",
    "BATCH_SIZE = 16\n",
    "EDGE_CROP = 16\n",
    "NB_EPOCHS = 5\n",
    "GAUSSIAN_NOISE = 0.1\n",
    "UPSAMPLE_MODE = 'SIMPLE'\n",
    "# downsampling inside the network\n",
    "NET_SCALING = None\n",
    "# downsampling in preprocessing\n",
    "IMG_SCALING = (1, 1)\n",
    "# number of validation images to use\n",
    "VALID_IMG_COUNT = 400\n",
    "# maximum number of steps_per_epoch in training\n",
    "MAX_TRAIN_STEPS = 200\n",
    "AUGMENT_BRIGHTNESS = False\n",
    "N_SAMPLE = 10\n",
    "IMG_SHAPE = (768, 768)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9edebd29",
   "metadata": {},
   "outputs": [],
   "source": [
    "hyperparams_gen = components.hyperparameters_gen(\n",
    "    num_epochs=5,\n",
    "    learning_rate=0.001,\n",
    "    batch_size=512,\n",
    "    hidden_units='64,64',\n",
    ")\n",
    "\n",
    "context.run(hyperparams_gen, enable_cache=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06bffdc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "json.load(\n",
    "    tf.io.gfile.GFile(\n",
    "        os.path.join(\n",
    "            hyperparams_gen.outputs['hyperparameters'].get()[0].uri, 'hyperparameters.json')\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2fd431a",
   "metadata": {},
   "source": [
    "## Model training"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "681153e5-88ad-45ba-8c7b-bfdc12d95e63",
   "metadata": {},
   "source": [
    "### Model Training Files"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "179261c9-5a82-4d30-9d2f-c1ff95543eef",
   "metadata": {},
   "source": [
    "#### trainer.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c5901fb-ed23-44dd-afea-5adb740b7a94",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting ./src/model_training/trainer.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile ./src/model_training/trainer.py\n",
    "import logging\n",
    "import tensorflow as tf\n",
    "import tensorflow.keras.backend as K\n",
    "from tensorflow.keras.losses import binary_crossentropy\n",
    "\n",
    "from tensorflow import keras\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "# from src.model_training \n",
    "import data, model\n",
    "from tensorflow.keras.callbacks import ModelCheckpoint, LearningRateScheduler, EarlyStopping, ReduceLROnPlateau\n",
    "\n",
    "def train(\n",
    "    train_data_dir,\n",
    "    eval_data_dir,\n",
    "    hyperparams,\n",
    "    base_model_dir=None,\n",
    "):\n",
    "    logging.basicConfig(level=logging.INFO)\n",
    "    logging.info(f\"Loading dataset from {train_data_dir}\")\n",
    "\n",
    "    train_dataset = data.get_dataset(\n",
    "        train_data_dir,\n",
    "        hyperparams,\n",
    "    )\n",
    "    \n",
    "    eval_dataset = data.get_dataset(\n",
    "        eval_data_dir,\n",
    "        hyperparams,\n",
    "    )\n",
    "    \n",
    "    def dice_coef(y_true, y_pred, smooth=1):\n",
    "        intersection = K.sum(y_true * y_pred, axis=[1,2,3])\n",
    "        union = K.sum(y_true, axis=[1,2,3]) + K.sum(y_pred, axis=[1,2,3])\n",
    "        return K.mean( (2. * intersection + smooth) / (union + smooth), axis=0)\n",
    "\n",
    "    def dice_p_bce(in_gt, in_pred):\n",
    "        return 1e-3*binary_crossentropy(in_gt, in_pred) - dice_coef(in_gt, in_pred)\n",
    "\n",
    "    def true_positive_rate(y_true, y_pred):\n",
    "        return K.sum(K.flatten(y_true)*K.flatten(K.round(y_pred)))/K.sum(y_true)\n",
    "    \n",
    "    \n",
    "    weight_path=\"{}_weights.best.ctph\".format('seg_model')\n",
    "\n",
    "    checkpoint = ModelCheckpoint(\n",
    "        weight_path, \n",
    "        monitor='val_dice_coef', \n",
    "        verbose=1, \n",
    "        save_best_only=True, \n",
    "        mode='max', \n",
    "        save_weights_only = True)\n",
    "\n",
    "    optimizer = keras.optimizers.Adam(\n",
    "        learning_rate=hyperparams[\"learning_rate\"], \n",
    "        decay=hyperparams[\"decay_rate\"])\n",
    "    \n",
    "    metrics = [dice_coef, 'binary_accuracy', true_positive_rate]\n",
    "    \n",
    "    reduceLROnPlat = ReduceLROnPlateau(\n",
    "        monitor='val_dice_coef', \n",
    "        factor=0.5, \n",
    "        patience=10,                       \n",
    "        verbose=1, \n",
    "        mode='max', \n",
    "        epsilon=0.0001, \n",
    "        cooldown=2, \n",
    "        min_lr=1e-6\n",
    "    )\n",
    "    \n",
    "    early_stopping = EarlyStopping(\n",
    "        monitor=\"val_dice_coef\", \n",
    "        mode=\"max\", \n",
    "        patience=30)\n",
    "    \n",
    "    callbacks_list = [checkpoint, reduceLROnPlat, early_stopping]\n",
    "\n",
    "    seg_model = model.create_model(hyperparams)\n",
    "    if base_model_dir:\n",
    "        try:\n",
    "            seg_model = keras.load_model(base_model_dir)\n",
    "        except:\n",
    "            pass\n",
    "\n",
    "    seg_model.compile(\n",
    "        optimizer = optimizer, \n",
    "        loss = dice_p_bce, \n",
    "        metrics = metrics\n",
    "    )\n",
    "\n",
    "    logging.info(\"Model training started...\")\n",
    "    loss_history = [\n",
    "        seg_model.fit(\n",
    "            train_dataset,\n",
    "            epochs=hyperparams[\"num_epochs\"], \n",
    "            validation_data=eval_dataset.take(1),\n",
    "            callbacks=callbacks_list,\n",
    "            verbose=1,\n",
    "            workers=1 # the generator is not very thread safe\n",
    "        )\n",
    "    ]\n",
    "    \n",
    "    logging.info(\"Model training completed.\")\n",
    "\n",
    "    return seg_model\n",
    "\n",
    "def evaluate(model, data_dir, raw_schema_location, hyperparams):\n",
    "    pass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a36e99c-c254-4635-8756-c7f7f3bac391",
   "metadata": {},
   "source": [
    "#### model.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59459325-17b7-4f6e-90e7-10db41fd5b39",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting ./src/model_training/model.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile ./src/model_training/model.py\n",
    "\"\"\"A DNN keras classification model.\"\"\"\n",
    "\n",
    "from keras import models, layers\n",
    "\n",
    "# Build U-Net model\n",
    "def upsample_conv(filters, kernel_size, strides, padding):\n",
    "    return layers.Conv2DTranspose(filters, kernel_size, strides=strides, padding=padding)\n",
    "\n",
    "def upsample_simple(filters, kernel_size, strides, padding):\n",
    "    return layers.UpSampling2D(strides)\n",
    "\n",
    "def UNet_keras(gaussian_noise=0.1, upsample_mode='SIMPLE', net_scaling=None, img_shape=(768, 768), edge_crop=16):\n",
    "    \n",
    "    input_img = layers.Input((img_shape[0], img_shape[1], 3), name = 'RGB_Input')\n",
    "    \n",
    "    pp_in_layer = input_img\n",
    "    \n",
    "    if upsample_mode=='DECONV':\n",
    "        upsample=upsample_conv\n",
    "    else:\n",
    "        upsample=upsample_simple\n",
    "    \n",
    "    if net_scaling is not None:\n",
    "        pp_in_layer = layers.AvgPool2D(net_scaling)(pp_in_layer)\n",
    "\n",
    "    pp_in_layer = layers.GaussianNoise(gaussian_noise)(pp_in_layer)\n",
    "    pp_in_layer = layers.BatchNormalization()(pp_in_layer)\n",
    "\n",
    "    c1 = layers.Conv2D(8, (3, 3), activation='relu', padding='same') (pp_in_layer)\n",
    "    c1 = layers.Conv2D(8, (3, 3), activation='relu', padding='same') (c1)\n",
    "    p1 = layers.MaxPooling2D((2, 2)) (c1)\n",
    "\n",
    "    c2 = layers.Conv2D(16, (3, 3), activation='relu', padding='same') (p1)\n",
    "    c2 = layers.Conv2D(16, (3, 3), activation='relu', padding='same') (c2)\n",
    "    p2 = layers.MaxPooling2D((2, 2)) (c2)\n",
    "\n",
    "    c3 = layers.Conv2D(32, (3, 3), activation='relu', padding='same') (p2)\n",
    "    c3 = layers.Conv2D(32, (3, 3), activation='relu', padding='same') (c3)\n",
    "    p3 = layers.MaxPooling2D((2, 2)) (c3)\n",
    "\n",
    "    c4 = layers.Conv2D(64, (3, 3), activation='relu', padding='same') (p3)\n",
    "    c4 = layers.Conv2D(64, (3, 3), activation='relu', padding='same') (c4)\n",
    "    p4 = layers.MaxPooling2D(pool_size=(2, 2)) (c4)\n",
    "\n",
    "    c5 = layers.Conv2D(128, (3, 3), activation='relu', padding='same') (p4)\n",
    "    c5 = layers.Conv2D(128, (3, 3), activation='relu', padding='same') (c5)\n",
    "\n",
    "    u6 = upsample(64, (2, 2), strides=(2, 2), padding='same') (c5)\n",
    "    u6 = layers.concatenate([u6, c4])\n",
    "    c6 = layers.Conv2D(64, (3, 3), activation='relu', padding='same') (u6)\n",
    "    c6 = layers.Conv2D(64, (3, 3), activation='relu', padding='same') (c6)\n",
    "\n",
    "    u7 = upsample(32, (2, 2), strides=(2, 2), padding='same') (c6)\n",
    "    u7 = layers.concatenate([u7, c3])\n",
    "    c7 = layers.Conv2D(32, (3, 3), activation='relu', padding='same') (u7)\n",
    "    c7 = layers.Conv2D(32, (3, 3), activation='relu', padding='same') (c7)\n",
    "\n",
    "    u8 = upsample(16, (2, 2), strides=(2, 2), padding='same') (c7)\n",
    "    u8 = layers.concatenate([u8, c2])\n",
    "    c8 = layers.Conv2D(16, (3, 3), activation='relu', padding='same') (u8)\n",
    "    c8 = layers.Conv2D(16, (3, 3), activation='relu', padding='same') (c8)\n",
    "\n",
    "    u9 = upsample(8, (2, 2), strides=(2, 2), padding='same') (c8)\n",
    "    u9 = layers.concatenate([u9, c1], axis=3)\n",
    "    c9 = layers.Conv2D(8, (3, 3), activation='relu', padding='same') (u9)\n",
    "    c9 = layers.Conv2D(8, (3, 3), activation='relu', padding='same') (c9)\n",
    "\n",
    "    d = layers.Conv2D(1, (1, 1), activation='sigmoid') (c9)\n",
    "    d = layers.Cropping2D((edge_crop, edge_crop))(d)\n",
    "    d = layers.ZeroPadding2D((edge_crop, edge_crop))(d)\n",
    "    \n",
    "    if net_scaling is not None:\n",
    "        d = layers.UpSampling2D(net_scaling)(d)\n",
    "\n",
    "    seg_model = models.Model(inputs=[input_img], outputs=[d])\n",
    "    \n",
    "    return seg_model\n",
    "\n",
    "def create_model(hyperparams):\n",
    "    return UNet_keras(\n",
    "        gaussian_noise=hyperparams[\"guassian_noise\"], \n",
    "        upsample_mode=hyperparams[\"upsample_mode\"], \n",
    "        net_scaling=hyperparams[\"net_scaling\"], \n",
    "        img_shape=hyperparams[\"img_shape\"], \n",
    "        edge_crop=hyperparams[\"edge_crop\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "424bdc02-cd75-4a93-95fd-19d2cf22d27d",
   "metadata": {},
   "source": [
    "#### data.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d43fe11-4caf-43db-8fda-a9a6765347d7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting ./src/model_training/data.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile ./src/model_training/data.py\n",
    "\"\"\"Functions for reading data as tf.data.Dataset.\"\"\"\n",
    "\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from tensorflow.keras import layers\n",
    "\n",
    "from google.cloud import storage\n",
    "\n",
    "class Augment(tf.keras.layers.Layer):\n",
    "    def __init__(self,  resize_shape=(768, 768), train=True, seed=42):\n",
    "        super().__init__()\n",
    "    # both use the same seed, so they'll make the same random changes.\n",
    "        seed = np.random.randint(1000)\n",
    "        if train:\n",
    "            self.augment_inputs = tf.keras.Sequential(\n",
    "                                    [\n",
    "                                        layers.experimental.preprocessing.RandomFlip(seed=seed),\n",
    "                                        layers.experimental.preprocessing.RandomRotation(0.1, seed=seed),\n",
    "                                        layers.experimental.preprocessing.RandomHeight(0.1, seed=seed),\n",
    "                                        layers.experimental.preprocessing.RandomWidth(0.1, seed=seed),\n",
    "                                        layers.experimental.preprocessing.RandomZoom(0.9, seed=seed),\n",
    "                                        layers.experimental.preprocessing.Rescaling(1.0 / 255),\n",
    "                                        layers.experimental.preprocessing.Resizing(resize_shape[0], resize_shape[0])\n",
    "                                    ]\n",
    "                                )\n",
    "\n",
    "            self.augment_labels = tf.keras.Sequential(\n",
    "                                    [\n",
    "                                        layers.experimental.preprocessing.RandomFlip(seed=seed),\n",
    "                                        layers.experimental.preprocessing.RandomRotation(0.1, seed=seed),\n",
    "                                        layers.experimental.preprocessing.RandomHeight(0.1, seed=seed),\n",
    "                                        layers.experimental.preprocessing.RandomWidth(0.1, seed=seed),\n",
    "                                        layers.experimental.preprocessing.RandomZoom(0.9, seed=seed),\n",
    "                                        layers.experimental.preprocessing.Resizing(resize_shape[0], resize_shape[0])\n",
    "                                    ]\n",
    "                                )\n",
    "        else:\n",
    "            self.augment_inputs = tf.keras.Sequential(\n",
    "                                    [\n",
    "                                        layers.experimental.preprocessing.Rescaling(1.0 / 255),\n",
    "                                        layers.experimental.preprocessing.Resizing(resize_shape[0], resize_shape[0])\n",
    "                                    ]\n",
    "                                )\n",
    "\n",
    "            self.augment_labels = tf.keras.Sequential(\n",
    "                                    [\n",
    "                                        layers.experimental.preprocessing.Resizing(resize_shape[0], resize_shape[0])\n",
    "                                    ]\n",
    "                                )       \n",
    "\n",
    "    def call(self, inputs, labels):\n",
    "        inputs = self.augment_inputs(inputs)\n",
    "        labels = self.augment_labels(labels)\n",
    "        return inputs, labels\n",
    "\n",
    "\n",
    "def rle_decode_tf(mask_rle, shape=(768, 768)):\n",
    "\n",
    "    shape = tf.convert_to_tensor(shape, tf.int64)\n",
    "    size = tf.math.reduce_prod(shape)\n",
    "    # Split string\n",
    "    s = tf.strings.split(mask_rle)\n",
    "    s = tf.strings.to_number(s, tf.int64)\n",
    "    # Get starts and lengths\n",
    "    starts = s[::2] - 1\n",
    "    lens = s[1::2]\n",
    "    # Make ones to be scattered\n",
    "    total_ones = tf.reduce_sum(lens)\n",
    "    ones = tf.ones([total_ones], tf.uint8)\n",
    "    # Make scattering indices\n",
    "    r = tf.range(total_ones)\n",
    "    lens_cum = tf.math.cumsum(lens)\n",
    "    s = tf.searchsorted(lens_cum, r, 'right')\n",
    "    idx = r + tf.gather(starts - tf.pad(lens_cum[:-1], [(1, 0)]), s)\n",
    "    # Scatter ones into flattened mask\n",
    "    mask_flat = tf.scatter_nd(tf.expand_dims(idx, 1), ones, [size])\n",
    "    return tf.expand_dims(tf.transpose(tf.reshape(mask_flat, shape)), axis=2)\n",
    "\n",
    "def multi_rle_encode(img):\n",
    "    labels = label(img[:, :, 0])\n",
    "    return [rle_encode(labels==k) for k in np.unique(labels[labels>0])]\n",
    "\n",
    "# ref: https://www.kaggle.com/paulorzp/run-length-encode-and-decode\n",
    "def rle_encode(img):\n",
    "    '''\n",
    "    img: numpy array, 1 - mask, 0 - background\n",
    "    Returns run length as string formated\n",
    "    '''\n",
    "    pixels = img.T.flatten()\n",
    "    pixels = np.concatenate([[0], pixels, [0]])\n",
    "    runs = np.where(pixels[1:] != pixels[:-1])[0] + 1\n",
    "    runs[1::2] -= runs[::2]\n",
    "    return ' '.join(str(x) for x in runs)\n",
    "\n",
    "def rle_decode(mask_rle, shape=(768, 768)):\n",
    "    '''\n",
    "    mask_rle: run-length as string formated (start length)\n",
    "    shape: (height,width) of array to return \n",
    "    Returns numpy array, 1 - mask, 0 - background\n",
    "    '''\n",
    "    s = mask_rle.split()\n",
    "    starts, lengths = [np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])]\n",
    "    starts -= 1\n",
    "    ends = starts + lengths\n",
    "    img = np.zeros(shape[0]*shape[1], dtype=np.uint8)\n",
    "    for lo, hi in zip(starts, ends):\n",
    "        img[lo:hi] = 1\n",
    "    return img.reshape(shape).T   # Needed to align to RLE direction\n",
    "\n",
    "def masks_as_image(in_mask_list):\n",
    "    #in_mask_list = tf.compat.as_str_any(in_mask_list)\n",
    "    # Take the individual ship masks and create a single mask array for all ships\n",
    "    all_masks = np.zeros((768, 768), dtype = np.int16)\n",
    "    #if isinstance(in_mask_list, list):\n",
    "    for mask in in_mask_list:\n",
    "        if isinstance(mask, str):\n",
    "            all_masks += rle_decode(mask)\n",
    "    return np.expand_dims(all_masks, -1)\n",
    "\n",
    "def merge_rle_encode(mask_rle, shape=(768, 768)):\n",
    "    img = np.zeros(shape=shape, dtype=np.uint8)\n",
    "\n",
    "    for rle in mask_rle.split(\";\"):\n",
    "        img += rle_decode(rle)\n",
    "\n",
    "    return rle_encode(img)\n",
    "\n",
    "def parse_db_to_img(filename, label):\n",
    "    file_path = filename\n",
    "    img = tf.io.read_file(file_path)\n",
    "    image = tf.image.decode_jpeg(img, channels=3)\n",
    "    label_img = rle_decode_tf(label)\n",
    "    return image, label_img\n",
    "\n",
    "def get_dataset(file_name, feature_spec):\n",
    "    \"\"\"Generates features and label for tuning/training.\n",
    "    Args:\n",
    "      file_pattern: input file path\n",
    "      feature_spec: a dictionary of feature specifications.\n",
    "      batch_size: representing the number of consecutive elements of returned\n",
    "        dataset to combine in a single batch\n",
    "    Returns:\n",
    "      A dataset that contains (features, indices) tuple where features is a\n",
    "        dictionary of Tensors, and indices is a single Tensor of label indices.\n",
    "    \"\"\"\n",
    "    \n",
    "    GCS_IMAGES = feature_spec['gcs_image']\n",
    "    IMG_SHAPE = feature_spec['img_shape']\n",
    "    GCS_BUCKET = feature_spec['gcs_bucket']\n",
    "    BATCH_SIZE = feature_spec['batch_size']\n",
    "    \n",
    "    bucket = storage.Client().bucket(GCS_BUCKET)\n",
    "    blob = bucket.blob(file_name)\n",
    "    blob.download_to_filename(file_name)\n",
    "    \n",
    "    train_df = pd.read_parquet(file_name)\n",
    "    dataset = tf.data.Dataset.from_tensor_slices((train_df['ImageId'].values, train_df['EncodedPixels'].values))\n",
    "    dataset = dataset.shuffle(buffer_size=100)\n",
    "    dataset = dataset.map(lambda x, y: parse_db_to_img(GCS_IMAGES + x, y))\n",
    "    dataset = dataset.batch(BATCH_SIZE)\n",
    "    dataset = dataset.map(Augment(resize_shape=IMG_SHAPE, train=True))\n",
    "    dataset = dataset.cache().prefetch(buffer_size=tf.data.AUTOTUNE)\n",
    "    return dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c3d2520-4cad-47a3-9fe4-4abc5b605c46",
   "metadata": {},
   "source": [
    "#### __ init __.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb6ad1db-12a8-419f-9454-a2609f25c4b2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "UsageError: %%writefile is a cell magic, but the cell body is empty.\n"
     ]
    }
   ],
   "source": [
    "%%writefile ./src/model_training/__init__.py\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e95ac0d8-af96-46cc-bec1-631097b220e5",
   "metadata": {},
   "source": [
    "#### defaults.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1da0bc1d-6fbf-4291-bdda-4ede4b98516e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting ./src/model_training/defaults.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile ./src/model_training/defaults.py\n",
    "\"\"\"Defaults for the model.\n",
    "These values can be tweaked to affect model training performance.\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "NET_SCALING = None\n",
    "IMG_SHAPE = (128, 128)\n",
    "GUASSIAN_NOISE = 0.1\n",
    "BATCH_SIZE = 16\n",
    "NUM_EPOCHS = 1\n",
    "NUM_EVAL_STEPS = 10\n",
    "EDGE_CROP = 16\n",
    "UPSAMPLE_MODE = \"SIMPLE\"\n",
    "LEARNING_RATE = 1e-4 \n",
    "DECAY_RATE = 1e-6\n",
    "\n",
    "GCS_IMAGES = \"gs://mle_airbus_dataset/train_v2/\"\n",
    "GCS_BUCKET = \"mle_airbus_dataset/\"\n",
    "\n",
    "def update_hyperparams(hyperparams: dict) -> dict:\n",
    "    if \"net_scaling\" not in hyperparams:\n",
    "        hyperparams[\"net_scaling\"] = NET_SCALING\n",
    "    if \"img_shape\" not in hyperparams:\n",
    "        hyperparams[\"img_shape\"] = IMG_SHAPE\n",
    "    if \"guassian_noise\" not in hyperparams:\n",
    "        hyperparams[\"num_epochs\"] = GUASSIAN_NOISE\n",
    "    if \"batch_size\" not in hyperparams:\n",
    "        hyperparams[\"batch_size\"] = BATCH_SIZE\n",
    "    if \"num_epochs\" not in hyperparams:\n",
    "        hyperparams[\"num_epochs\"] = NUM_EPOCHS\n",
    "    if \"num_eval_steps\" not in hyperparams:\n",
    "        hyperparams[\"num_eval_steps\"] = NUM_EVAL_STEPS\n",
    "    if \"edge_crop\" not in hyperparams:\n",
    "        hyperparams[\"edge_crop\"] = EDGE_CROP\n",
    "    if \"upsample_mode\" not in hyperparams:\n",
    "        hyperparams[\"upsample_mode\"] = UPSAMPLE_MODE\n",
    "    if \"learning_rate\" not in hyperparams:\n",
    "        hyperparams[\"learning_rate\"] = LEARNING_RATE\n",
    "    if \"decay_rate\" not in hyperparams:\n",
    "        hyperparams[\"decay_rate\"] = DECAY\n",
    "    \n",
    "    if \"gcs_image\" not in hyperparams:\n",
    "        hyperparams[\"gcs_image\"] = GCS_IMAGES\n",
    "    if \"gcs_bucket\" not in hyperparams:\n",
    "        hyperparams[\"gcs_bucket\"] = GCS_BUCKET\n",
    "    return hyperparams"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f77976b-de99-4a7c-b5c8-1c7edd4a151c",
   "metadata": {},
   "source": [
    "#### task.py\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0400336-973b-4ba2-b010-1d8e37799a3f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting ./src/model_training/task.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile ./src/model_training/task.py\n",
    "\"\"\"The entrypoint for the Vertex training job.\"\"\"\n",
    "\n",
    "import os\n",
    "import sys\n",
    "from datetime import datetime\n",
    "import logging\n",
    "import tensorflow as tf\n",
    "from tensorflow.python.client import device_lib\n",
    "import argparse\n",
    "from typing import Optional\n",
    "\n",
    "from google.cloud import aiplatform as vertex_ai\n",
    "\n",
    "# from src.model_training \n",
    "import defaults, trainer\n",
    "\n",
    "\n",
    "dirname = os.path.dirname(__file__)\n",
    "dirname = dirname.replace(\"/model_training\", \"\")\n",
    "\n",
    "\n",
    "def get_args():\n",
    "    parser = argparse.ArgumentParser()\n",
    "\n",
    "    parser.add_argument(\n",
    "        \"--model-dir\",\n",
    "        default=\"\",\n",
    "        type=str,\n",
    "    )\n",
    "\n",
    "    parser.add_argument(\n",
    "        \"--train-data-dir\",\n",
    "        type=str,\n",
    "    )\n",
    "\n",
    "    parser.add_argument(\n",
    "        \"--eval-data-dir\",\n",
    "        type=str,\n",
    "    )\n",
    "    \n",
    "    parser.add_argument(\n",
    "        \"--out-model\",\n",
    "        type=str,\n",
    "    )\n",
    "    \n",
    "    parser.add_argument(\"--net-scaling\", default=None, type=Optional[bool])\n",
    "    parser.add_argument(\"--image-shape\", default=(128,128), type=tuple)\n",
    "    parser.add_argument(\"--guassian-noise\", default=0.1, type=float)\n",
    "    parser.add_argument(\"--batch-size\", default=16, type=float)\n",
    "    parser.add_argument(\"--num-epochs\", default=1, type=int)\n",
    "    parser.add_argument(\"--num-eval-steps\", default=10, type=int)\n",
    "    parser.add_argument(\"--edge-crop\", default=16, type=int)\n",
    "    parser.add_argument(\"--upsample-mode\", default=\"SIMPLE\", type=str)\n",
    "    parser.add_argument(\"--learning-rate\", default=1e-4, type=float)\n",
    "    parser.add_argument(\"--decay-rate\", default=1e-6, type=float)\n",
    "\n",
    "    parser.add_argument(\"--gcs-bucket\", default=\"mle_airbus_dataset\", type=str)\n",
    "    parser.add_argument(\"--gcs-image\", default=\"gs://mle_airbus_dataset/train_v2/\", type=str)\n",
    "    \n",
    "    return parser.parse_args()\n",
    "\n",
    "\n",
    "def main() -> str:\n",
    "    args = get_args()\n",
    "    time = datetime.now()\n",
    "    hyperparams = vars(args)\n",
    "    hyperparams = defaults.update_hyperparams(hyperparams)\n",
    "    logging.info(f\"Hyperparameter: {hyperparams}\")\n",
    "\n",
    "#     if args.experiment_name:\n",
    "#         vertex_ai.init(\n",
    "#             project=args.project,\n",
    "#             staging_bucket=args.staging_bucket,\n",
    "#             experiment=args.experiment_name,\n",
    "#         )\n",
    "\n",
    "#         logging.info(f\"Using Vertex AI experiment: {args.experiment_name}\")\n",
    "\n",
    "#         run_id = args.run_name\n",
    "#         if not run_id:\n",
    "#             run_id = f\"run-gcp-{datetime.now().strftime('%Y%m%d%H%M%S')}\"\n",
    "\n",
    "#         vertex_ai.start_run(run_id)\n",
    "#         logging.info(f\"Run {run_id} started.\")\n",
    "\n",
    "#         vertex_ai.log_params(hyperparams)\n",
    "\n",
    "    seg_model = trainer.train(\n",
    "        train_data_dir=args.train_data_dir,\n",
    "        eval_data_dir=args.eval_data_dir,\n",
    "        hyperparams=hyperparams,\n",
    "    )\n",
    "\n",
    "    # val_loss, val_accuracy = trainer.evaluate(\n",
    "    #     model=classifier,\n",
    "    #     data_dir=args.eval_data_dir,\n",
    "    #     raw_schema_location=RAW_SCHEMA_LOCATION,\n",
    "    #     tft_output_dir=args.tft_output_dir,\n",
    "    #     hyperparams=hyperparams,\n",
    "    # )\n",
    "    \n",
    "    \n",
    "    # Report val_accuracy to Vertex hypertuner.\n",
    "    # logging.info(f'Reporting metric {HYPERTUNE_METRIC_NAME}={val_accuracy} to Vertex hypertuner...')\n",
    "    # hpt = hypertune.HyperTune()\n",
    "    # hpt.report_hyperparameter_tuning_metric(\n",
    "    #     hyperparameter_metric_tag=HYPERTUNE_METRIC_NAME,\n",
    "    #     metric_value=val_accuracy,\n",
    "    #     global_step=args.num_epochs * args.batch_size\n",
    "    # )\n",
    "\n",
    "    # Log metrics in Vertex Experiments.\n",
    "    # logging.info(f'Logging metrics to Vertex Experiments...')\n",
    "    # if args.experiment_name:\n",
    "    #     vertex_ai.log_metrics({\"val_loss\": val_loss, \"val_accuracy\": val_accuracy})\n",
    "    \n",
    "    \n",
    "    import pickle\n",
    "\n",
    "    logging.info(f\"exporting model\")\n",
    "    timestamp = time.strftime(\"%Y%m%d-%H%M%S\")\n",
    "\n",
    "    export_dir = \"gs://mle_airbus_dataset/trained_model/segm_full_{}\".format(timestamp)\n",
    "    print('Exporting to {}'.format(export_dir))\n",
    "    tf.saved_model.save(seg_model, export_dir)\n",
    "\n",
    "    # in two lines of code\n",
    "    # with open(hyperparams['model_dir'] + f\"segm_full_{timestamp}/loss.pickle\", \"wb\") as f:\n",
    "    #     print('Exporting to {}/loss.pickle'.format(export_dir))\n",
    "    #     pickle.dump(loss_history[0].history, f)\n",
    "\n",
    "    logging.info(f\"exported model: {export_dir}\")\n",
    "    # model_filename = args.model_dir + \"model_\" + datetime.datetime.now().strftime('%Y%m%d_%H%M')\n",
    "    # model.save_model(model_filename)\n",
    "\n",
    "    with open(hyperparams['out_model'], 'w') as f:\n",
    "        f.write(export_dir)\n",
    "    # os.stdout(export_dir)\n",
    "    return (print(export_dir))\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    logging.getLogger().setLevel(logging.INFO)\n",
    "    logging.info(f\"Python Version = {sys.version}\")\n",
    "    logging.info(f\"TensorFlow Version = {tf.__version__}\")\n",
    "    logging.info(f'TF_CONFIG = {os.environ.get(\"TF_CONFIG\", \"Not found\")}')\n",
    "    logging.info(f\"DEVICES = {device_lib.list_local_devices()}\")\n",
    "    logging.info(f\"Task started...\")\n",
    "    main()\n",
    "    logging.info(f\"Task completed.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79a72326-347e-4495-9a87-6ee437ef8b07",
   "metadata": {},
   "source": [
    "#### Create DockerFile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "023a5f08-9d1b-4bbd-95b6-fa7f66fbb376",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing Dockerfile\n"
     ]
    }
   ],
   "source": [
    "%%writefile Dockerfile\n",
    "FROM asia-docker.pkg.dev/vertex-ai/training/tf-cpu.2-8:latest\n",
    "WORKDIR /app\n",
    "COPY src/model_training src/model_training\n",
    "ENTRYPOINT [\"python\",\"./src/model_training/task.py\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ebec54d-f2e8-47f7-b6ca-d6171a7f00c0",
   "metadata": {},
   "source": [
    "#### Create Component"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0056b3c1-7806-40f1-9c54-779a6f2c6135",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Required Parameters, change according to your setups\n",
    "PROJECT_ID='mle-airbus-detection-smu'\n",
    "GCS_BUCKET='mle_airbus_dataset'\n",
    "REGION = 'asia-east1'\n",
    "ARTIFACT_REGISTRY_REPO=\"airbus-mle\"\n",
    "CONTAINER_REGISTRY=f\"{REGION}-docker.pkg.dev/{PROJECT_ID}/{ARTIFACT_REGISTRY_REPO}/trainer-airbus-model-ey:latest\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f123bc42-1ff6-4d98-a757-d3913d79af89",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "asia-east1-docker.pkg.dev/mle-airbus-detection-smu/airbus-mle/trainer-airbus-model-ey:latest\n"
     ]
    }
   ],
   "source": [
    "%%bash -s \"{CONTAINER_REGISTRY}\"\n",
    "\n",
    "CONTAINER_REGISTRY=$1\n",
    "echo ${CONTAINER_REGISTRY}\n",
    "# Create the component definition .yaml file\n",
    "cat > ./build/tensorflow_airbus.yaml <<HERE\n",
    "name: Tensorflow Airbus Model Training\n",
    "description: Train a tf model and save to GCS\n",
    "inputs:\n",
    "  - name: model_dir\n",
    "    description: 'Path to save model.'\n",
    "    type: String\n",
    "  - name: train_data_dir\n",
    "    description: 'Training dataset directory.'\n",
    "    type: String\n",
    "  - name: eval_data_dir\n",
    "    description: 'Evaluation dataset directory.'\n",
    "    type: String\n",
    "    \n",
    "outputs:\n",
    "  - name: gcs_model_path\n",
    "    description: 'Trained model path.'\n",
    "    type: String\n",
    "implementation:\n",
    "    container:\n",
    "        image: ${CONTAINER_REGISTRY}\n",
    "        command: [\n",
    "          python, ./src/model_training/task.py,\n",
    "          --model-dir, {inputValue: model_dir},\n",
    "          --train-data-dir, {inputValue: train_data_dir},\n",
    "          --eval-data-dir, {inputValue: eval_data_dir},\n",
    "          --out-model: {outputPath: gcs_model_path}\n",
    "        ]\n",
    "HERE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15166d0b-fae6-4c5a-b69b-4f61435c9f02",
   "metadata": {},
   "outputs": [],
   "source": [
    "import kfp\n",
    "trainer_component = kfp.components.load_component_from_file(\"./build/tensorflow_airbus.yaml\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c023e9cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tfx.dsl.components.common.resolver import Resolver\n",
    "from tfx.dsl.experimental import latest_artifacts_resolver\n",
    "from tfx.dsl.experimental import latest_blessed_model_resolver"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab0b69e7",
   "metadata": {},
   "source": [
    "### Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3606d40e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import subprocess\n",
    "cmd = ['docker', 'build', '-f', 'Dockerfile_trainer', '.', '--tag', CONTAINER_REGISTRY]\n",
    "\n",
    "build_log = (subprocess.run(cmd, stdout=subprocess.PIPE).stdout[:-1].decode('utf-8'))\n",
    "print(build_log)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a9c09c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import subprocess\n",
    "cmd = ['docker', 'run', CONTAINER_REGISTRY, \n",
    "       '--model-dir', f\"gs://{GCS_BUCKET}/airbusmlepipeline/{MODEL_OUTPUT_NAME}\", \n",
    "       '--train-data-dir', f\"train.parquet\", \n",
    "       '--eval-data-dir', f\"test.parquet\"]\n",
    "\n",
    "build_log = (subprocess.run(cmd, stdout=subprocess.PIPE).stdout[:-1].decode('utf-8'))\n",
    "print(build_log)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "766a7478",
   "metadata": {},
   "source": [
    "### Evaluate and validate the model against the baseline model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3e939ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluation_results = evaluator.outputs['evaluation'].get()[0].uri\n",
    "print(\"validation_ok:\", tfma.load_validation_result(evaluation_results).validation_ok, '\\n')\n",
    "\n",
    "for entry in list(tfma.load_metrics(evaluation_results))[0].metric_keys_and_values:\n",
    "    value = entry.value.double_value.value\n",
    "    if value:\n",
    "        print(entry.key.name, \":\", round(entry.value.double_value.value, 3))"
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "python3",
   "name": "tf2-gpu.2-8.m93",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/tf2-gpu.2-8:m93"
  },
  "interpreter": {
   "hash": "447788c8e7023b8ad5a7a5b6695c585102448263744d543b03e0f5f55d0a66b9"
  },
  "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.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}