full_palette_generation/notebooks/train_model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "41cb4ea0",
   "metadata": {},
   "source": [
    "### This is the pre-training of text-aware masked color model for full palette generation\n",
    "- It takes 10mins on GPU (Tesla T4 * 1) for one time training with early stop (patience=30)\n",
    "- For a quick start, do Pretraining with pre-created color corpus files and text embedding files\n",
    "    - To create color and text data for training, please check preprocess.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7a5b9029",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-09-04 05:31:20.996437: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcudart.so.11.0\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-09-04 05:31:22.732604: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcuda.so.1\n",
      "2023-09-04 05:31:22.740454: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-09-04 05:31:22.741027: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1716] Found device 0 with properties: \n",
      "pciBusID: 0000:00:04.0 name: Tesla T4 computeCapability: 7.5\n",
      "coreClock: 1.59GHz coreCount: 40 deviceMemorySize: 14.75GiB deviceMemoryBandwidth: 298.08GiB/s\n",
      "2023-09-04 05:31:22.741094: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcudart.so.11.0\n",
      "2023-09-04 05:31:22.745369: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcublas.so.11\n",
      "2023-09-04 05:31:22.747363: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcufft.so.10\n",
      "2023-09-04 05:31:22.747744: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcurand.so.10\n",
      "2023-09-04 05:31:22.751155: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcusolver.so.10\n",
      "2023-09-04 05:31:22.752279: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcusparse.so.11\n",
      "2023-09-04 05:31:22.752537: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcudnn.so.8\n",
      "2023-09-04 05:31:22.752696: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-09-04 05:31:22.753460: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:982] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-09-04 05:31:22.754048: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1858] Adding visible gpu devices: 0\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "\n",
    "print(len(tf.config.experimental.list_physical_devices('GPU')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a7537cff",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from collections import defaultdict  # For word frequency\n",
    "import os\n",
    "import math\n",
    "import random\n",
    "import ast\n",
    "from datetime import datetime\n",
    "from collections import Counter\n",
    "import numpy as np\n",
    "\n",
    "import sys\n",
    "sys.path.append('../src')\n",
    "\n",
    "import full_palette_generation.text_color_model.color_model as Model\n",
    "from full_palette_generation.text_color_model.input_data_generator import DataGenerator\n",
    "from full_palette_generation.text_color_model.model_config import Config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "80507f36",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_pretrain_task_accuracy(mlm_predict, batch_mlm_mask, origin_x):\n",
    "    \n",
    "    batch_mlm_mask = tf.cast(batch_mlm_mask, dtype=tf.int32)\n",
    "    index = tf.where(batch_mlm_mask == 1)\n",
    "    x_predict = tf.math.argmax(mlm_predict, axis=-1)\n",
    "    x_predict = tf.gather_nd(x_predict, index)\n",
    "    x_real = tf.gather_nd(origin_x, index)\n",
    "    mlm_accuracy = tf.keras.metrics.Accuracy()\n",
    "    mlm_accuracy.update_state(x_predict, x_real)\n",
    "    mlm_accuracy = mlm_accuracy.result().numpy()\n",
    "\n",
    "    return mlm_accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d35ab4aa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Layer text2_palettes_3 is casting an input tensor from dtype float64 to the layer's dtype of float32, which is new behavior in TensorFlow 2.  The layer has dtype float32 because its dtype defaults to floatx.\n",
      "\n",
      "If you intended to run this layer in float32, you can safely ignore this warning. If in doubt, this warning is likely only an issue if you are porting a TensorFlow 1.X model to TensorFlow 2.\n",
      "\n",
      "To change all layers to have dtype float64 by default, call `tf.keras.backend.set_floatx('float64')`. To change just this layer, pass dtype='float64' to the layer constructor. If you are the author of this layer, you can disable autocasting by passing autocast=False to the base Layer constructor.\n",
      "\n",
      "WARNING:tensorflow:Layer text2_palettes__loss_3 is casting an input tensor from dtype float64 to the layer's dtype of float32, which is new behavior in TensorFlow 2.  The layer has dtype float32 because its dtype defaults to floatx.\n",
      "\n",
      "If you intended to run this layer in float32, you can safely ignore this warning. If in doubt, this warning is likely only an issue if you are porting a TensorFlow 1.X model to TensorFlow 2.\n",
      "\n",
      "To change all layers to have dtype float64 by default, call `tf.keras.backend.set_floatx('float64')`. To change just this layer, pass dtype='float64' to the layer constructor. If you are the author of this layer, you can disable autocasting by passing autocast=False to the base Layer constructor.\n",
      "\n",
      "Epoch 0, step 7, loss 2.5326, mlm_loss 2.5326, mlm_acc 0.0246\n",
      "Val: Epoch 0, step 0, loss 2.4810, mlm_loss 2.4810, mlm_acc 0.0321\n",
      "Test: Epoch 1, step 0, loss 2.4660, mlm_loss 2.4660, mlm_acc 0.0283\n"
     ]
    }
   ],
   "source": [
    "# pretrain\n",
    "\n",
    "for n in range(1):\n",
    "    # physical_devices = tf.config.experimental.list_physical_devices('CPU')\n",
    "    physical_devices = tf.config.experimental.list_physical_devices('GPU')\n",
    "    assert len(physical_devices) > 0, \"Not enough GPU hardware devices available\"\n",
    "    tf.config.experimental.set_memory_growth(physical_devices[0], True)\n",
    "\n",
    "    model = Model.Text2Palettes(Config)\n",
    "    optimizer = tf.keras.optimizers.Adam(Config['Learning_Rate'])\n",
    "    loss_fn = Model.Text2Palettes_Loss()\n",
    "    dataset = DataGenerator(Config)\n",
    "\n",
    "    # create the data for validation and test\n",
    "    PROJECT_PATH = Config['project_path']\n",
    "    representation = Config['representation']\n",
    "    text_model = Config['text_model']\n",
    "    emb_file = Config['emb_file']\n",
    "\n",
    "    Config_val = Config.copy()\n",
    "    Config_val['Corpus_File_Path'] = os.path.join(PROJECT_PATH, f'color/color_corpus_{representation}_val.txt')\n",
    "    Config_val['Text_Input_File_Path'] = os.path.join(PROJECT_PATH, f'text/text_input_val.txt')\n",
    "    Config_val['Text_Input_Emb_File_Path'] = os.path.join(PROJECT_PATH, f'text/{emb_file}/text_input_emb{text_model}_val.txt')\n",
    "\n",
    "    Config_test = Config.copy()\n",
    "    Config_test['Corpus_File_Path'] = os.path.join(PROJECT_PATH, f'color/color_corpus_{representation}_test.txt')\n",
    "    Config_test['Text_Input_File_Path'] = os.path.join(PROJECT_PATH, f'text/text_input_test.txt')\n",
    "    Config_test['Text_Input_Emb_File_Path'] = os.path.join(PROJECT_PATH, f'text/{emb_file}/text_input_emb{text_model}_test.txt')\n",
    "\n",
    "    dataset_val = DataGenerator(Config_val)\n",
    "    dataset_test = DataGenerator(Config_test)\n",
    "\n",
    "\n",
    "    patience = 50 # baseline:10\n",
    "    best = math.inf\n",
    "    wait = 0\n",
    "\n",
    "    Config['Saved_Weight'] = os.path.join(PROJECT_PATH, f'Saved_Weight{text_model}_{Config[\"Embedding_Size\"]}d_{representation}_{Config[\"Mask_Rate\"]}_{Config[\"Mask_Token_Rate\"]}_v{n}')\n",
    "    checkpoint = tf.train.Checkpoint(model=model)\n",
    "    checkpoint.restore(tf.train.latest_checkpoint(Config['Saved_Weight']))\n",
    "    manager = tf.train.CheckpointManager(checkpoint, directory=Config['Saved_Weight'], max_to_keep=5)\n",
    "    log_dir = os.path.join(Config['Log_Dir'], datetime.now().strftime(\"%Y-%m-%d\"))\n",
    "    writer = tf.summary.create_file_writer(log_dir)\n",
    "    \n",
    "    EPOCH = 2000 # 2000 when training\n",
    "    for epoch in range(EPOCH):\n",
    "        for step in range(len(dataset)):\n",
    "            batch_x, batch_mlm_mask, batch_mcc_mask, origin_x, batch_segment, batch_padding_mask, batch_text_input_embed = dataset[step]\n",
    "\n",
    "            with tf.GradientTape() as t:\n",
    "                mlm_predict, sequence_output = model((batch_x, batch_mlm_mask, batch_segment, batch_text_input_embed), training=True)\n",
    "\n",
    "                mlm_loss = loss_fn((mlm_predict, batch_mlm_mask, origin_x))\n",
    "                mlm_loss = tf.reduce_mean(mlm_loss)\n",
    "\n",
    "                loss = mlm_loss\n",
    "\n",
    "            gradients = t.gradient(loss, model.trainable_variables)\n",
    "            optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
    "\n",
    "            mlm_acc = calculate_pretrain_task_accuracy(mlm_predict, batch_mlm_mask, origin_x)\n",
    "\n",
    "            if step == len(dataset) - 1 and epoch % 10 == 0:\n",
    "                print(\n",
    "                    'Epoch {}, step {}, loss {:.4f}, mlm_loss {:.4f}, mlm_acc {:.4f}'.format(\n",
    "                        epoch, step, loss.numpy(),\n",
    "                        mlm_loss.numpy(),\n",
    "                        mlm_acc,\n",
    "                        ))\n",
    "\n",
    "        for val_step in range(len(dataset_val)):\n",
    "            val_batch_x, val_batch_mlm_mask, val_batch_mcc_mask, val_origin_x, val_batch_segment, val_batch_padding_mask, val_batch_text_input_embed = dataset_val[val_step]\n",
    "\n",
    "            val_mlm_predict, val_sequence_output = model((val_batch_x, val_batch_mlm_mask, val_batch_segment, val_batch_text_input_embed), training=False)\n",
    "\n",
    "            val_mlm_loss = loss_fn((val_mlm_predict, val_batch_mlm_mask, val_origin_x))\n",
    "            val_mlm_loss = tf.reduce_mean(val_mlm_loss)\n",
    "            \n",
    "            val_mlm_acc = calculate_pretrain_task_accuracy(val_mlm_predict, val_batch_mlm_mask, val_origin_x)\n",
    "\n",
    "            val_loss = val_mlm_loss\n",
    "\n",
    "            if val_step == len(dataset_val) - 1 and epoch % 10 == 0:\n",
    "                print(\n",
    "                    'Val: Epoch {}, step {}, loss {:.4f}, mlm_loss {:.4f}, mlm_acc {:.4f}'.format(\n",
    "                        epoch, val_step, val_loss.numpy(),\n",
    "                        val_mlm_loss.numpy(),\n",
    "                        val_mlm_acc,\n",
    "                        ))\n",
    "        \n",
    "        path = manager.save(checkpoint_number=epoch)\n",
    "\n",
    "        # early stopping\n",
    "        wait += 1\n",
    "        if val_loss < best:\n",
    "            best = val_loss\n",
    "            wait = 0\n",
    "        if wait >= patience:\n",
    "            break\n",
    "            \n",
    "    for test_step in range(len(dataset_test)):\n",
    "        test_batch_x, test_batch_mlm_mask, test_batch_mcc_mask, test_origin_x, test_batch_segment, test_batch_padding_mask, test_batch_text_input_embed = dataset_test[test_step]\n",
    "            \n",
    "        test_mlm_predict, test_sequence_output = model((test_batch_x, test_batch_mlm_mask, test_batch_segment, test_batch_text_input_embed), training=False)\n",
    "\n",
    "        test_mlm_loss = loss_fn((test_mlm_predict, test_batch_mlm_mask, test_origin_x))\n",
    "        test_mlm_loss = tf.reduce_mean(test_mlm_loss)\n",
    "\n",
    "        test_mlm_acc = calculate_pretrain_task_accuracy(test_mlm_predict, test_batch_mlm_mask, test_origin_x)\n",
    "\n",
    "        test_loss = test_mlm_loss\n",
    "        \n",
    "        if test_step == len(dataset_test) - 1:\n",
    "            print(\n",
    "                'Test: Epoch {}, step {}, loss {:.4f}, mlm_loss {:.4f}, mlm_acc {:.4f}'.format(\n",
    "                    epoch, test_step, test_loss.numpy(),\n",
    "                    test_mlm_loss.numpy(),\n",
    "                    test_mlm_acc,\n",
    "                    ))\n",
    "    \n",
    "    # save model\n",
    "    model.save(f'../data/trained_model/t2p_ca1_mca1_stop{patience}_lr{Config[\"Learning_Rate\"]}_{text_model}_{Config[\"Embedding_Size\"]}d_{representation}_{Config[\"Mask_Rate\"]}_{Config[\"Mask_Token_Rate\"]}_{n}')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a35bc7f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "environment": {
   "name": "tf2-gpu.2-3.m79",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/tf2-gpu.2-3:m79"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}