demo/denas/hf/DENAS_HF_DEMO.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![open in colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/intel/e2eAIOK/blob/main/demo/denas/hf/DENAS_HF_DEMO.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# AIOK DE-NAS HF Demo\n",
    "\n",
    "DE-NAS is a multi-model, hardware-aware, train-free NAS to construct compact model architectures for target platform directly. DE-NAS includes CNN-based search space for CV domain and Transformer-based search space for CV/NLP/ASR domains, and leverages hardware-aware train-free scoring method to evaluate the performance of the candidate architecture without training.\n",
    "\n",
    "This demo mainly introduces HF integration with DE-NAS to search lighter, faster, higher performance transformer-based HF model in a training-free way."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Content\n",
    "\n",
    "* [Overview](#overview)\n",
    "    * [DE-NAS on HF Domain](#de-nas-on-hf-domain)\n",
    "* [Getting Started](#getting-started)\n",
    "    * [1. Enviroment Setup](#1-environment-setup)\n",
    "    * [2. Workflow Prepare](#2-workflow-prepare)\n",
    "    * [3. Data Prepare](#3-data-prepare)\n",
    "    * [4. Launch Search](#4-launch-search)\n",
    "    * [5. Launch Training with Best Searched Model Structure](#5-launch-training-with-best-searched-model-structure)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Overview\n",
    "\n",
    "## DE-NAS on HF Domain\n",
    "\n",
    "### DE-NAS on HF Search Space and Supernet\n",
    "Transformer-based search space consists of number of transformer layer, number of attention head, size of query/key/value, size of MLP, and dimension of embeddings, and the supernet of DE-NAS on HF is a HF BERT-based structure, which are shown as the below figure.\n",
    "\n",
    "<center>\n",
    "<img src=\"./img/HF_BERT_Search_Space.png\" width=\"800\"/><figure>DE-NAS on HF BERT search space</figure>\n",
    "</center>\n",
    "\n",
    "### DE-NAS Searched HF BERT Architecture\n",
    "By deploying the train-free EA search engine on DE-NAS HF BERT search space and supernet, the DE-NAS HF BERT delivered the architecture that was more compact than the HF BERT-Base model as shown in the below figure:\n",
    "\n",
    "<center>\n",
    "<img src=\"./img/DENAS HF BERT Architecture.png\" width=\"500\"/><figure>DE-NAS Searched HF BERT Architecture</figure>\n",
    "</center>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Getting Started\n",
    "Noted: Need to download dataset and pre-trained model manually to run this demo.\n",
    "\n",
    "## 1. Environment Setup\n",
    "\n",
    "### (Option 1) Use Pip install"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install e2eAIOK-denas --pre"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### (Option 2) Use Docker\n",
    "\n",
    "Step1. prepare code\n",
    "\n",
    "``` shell\n",
    "### Build docker image ###\n",
    "# clone the e2eaiok repo\n",
    "git clone https://github.com/intel/e2eAIOK.git\n",
    "cd e2eAIOK\n",
    "git submodule update --init --recursive\n",
    "```\n",
    "\n",
    "Step2. build docker image\n",
    "\n",
    "```shell\n",
    "python3 scripts/start_e2eaiok_docker.py -b pytorch112 -w ${host0} ${host1} ${host2} ${host3} --proxy \"\"\n",
    "```\n",
    "\n",
    "Step3. run docker and start conda env\n",
    "\n",
    "``` shell\n",
    "sshpass -p docker ssh ${host0} -p 12347\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Enter Docker\n",
    "\n",
    "``` shell\n",
    "# connect the docker\n",
    "sshpass -p docker ssh ${host0} -p 12347\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Workflow Prepare\n",
    "\n",
    "* Conf for HF BERT DE-NAS Search\n",
    "\n",
    "```yaml\n",
    "# conf for HF bert\n",
    "model_type: hf\n",
    "search_engine: EvolutionarySearchEngine\n",
    "pretrained_model_path: /home/vmagent/app/dataset/\n",
    "batch_size: 32\n",
    "\n",
    "# conf for evolutionary search engine\n",
    "random_max_epochs: 1000 #random search max epochs\n",
    "max_epochs: 10 #search epoch\n",
    "select_num: 50\n",
    "population_num: 50\n",
    "m_prob: 0.2\n",
    "s_prob: 0.4\n",
    "crossover_num: 25\n",
    "mutation_num: 25\n",
    "img_size: 128\n",
    "max_param_limits: 110\n",
    "min_param_limits: 55\n",
    "seed: 0\n",
    "\n",
    "# enable/disable each DE-Score\n",
    "expressivity_weight: 0\n",
    "complexity_weight: 0\n",
    "diversity_weight: 0.00001\n",
    "saliency_weight: 1\n",
    "latency_weight: 0.01\n",
    "```\n",
    "\n",
    "The above yaml-format file shows the DE-NAS search relevant configuration on BERT, which was placed on the `/home/vmagent/app/e2eaiok/conf/denas/hf/e2eaiok_denas_hf.conf`. It determines the type of search engine, search hyparameter (etc., batch_size, select_num and population_num), DE-Score parameters (etc., expressivity score weight and latency weight) and supernet/search space configuration (etc., supernet_cfg).\n",
    "\n",
    "* Conf for BERT Supernet and Search Space\n",
    "\n",
    "```yaml\n",
    "# HF BERT supernet definition\n",
    "supernet:\n",
    "  bert-base-uncased\n",
    "\n",
    "# BERT search space definition\n",
    "#search_space:\n",
    "#  intermediate_size: \n",
    "#    max: 3072\n",
    "#    step: 16\n",
    "```\n",
    "\n",
    "The above yaml-format file describes the details of BERT-base supernet and search space configuration, which was also placed on the `/home/vmagent/app/e2eaiok/conf/denas/hf/e2eaiok_denas_hf.conf`. The common HF layer structure of \"hidden_size\", \"num_attention_heads\" and \"num_hidden_layers\" are the default items used in the DE-NAS HF search space, and other structure items (etc., \"embedding_size\") are determined by the user-self.\n",
    "\n",
    "* Download pre-trained model from Hugging Face\n",
    "    * Download and extract one of BERT-Base-Uncased pretrained models from [Hugging Face repository](https://huggingface.co/bert-base-uncased/tree/main) to `/home/vmagent/app/dataset/bert-base-uncased/`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Data Prepare\n",
    "\n",
    "* Prepare Dataset\n",
    "    * Download Dataset: Stanford Question Answering Dataset (SQuAD) is a reading comprehension dataset, consisting of questions posed by crowdworkers on a set of Wikipedia articles, where the answer to every question is a segment of text, or span, from the corresponding reading passage, or the question might be unanswerable. SQuAD 1.1 contains 100,000+ question-answer pairs on 500+ articles.\n",
    "    * Download from below path to `/home/vmagent/app/dataset/SQuAD`\n",
    "        * Train Data: [train-v1.1.json](https://rajpurkar.github.io/SQuAD-explorer/dataset/train-v1.1.json)\n",
    "        * Test Data: [dev-v1.1.json](https://rajpurkar.github.io/SQuAD-explorer/dataset/dev-v1.1.json)\n",
    "``` bash\n",
    "Data Format:\n",
    "{\n",
    "    \"answers\": {\n",
    "        \"answer_start\": [1],\n",
    "        \"text\": [\"This is a test text\"]\n",
    "    },\n",
    "    \"context\": \"This is a test context.\",\n",
    "    \"id\": \"1\",\n",
    "    \"question\": \"Is this a test?\",\n",
    "    \"title\": \"train test\"\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Launch Search"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Launch DE-NAS search process on HF model BERT with the input of overall search configuration `/home/vmagent/app/e2eaiok/conf/denas/hf/e2eaiok_denas_hf.conf`, and will produce the best model structure as a tuple `(layer_num, head_num, hidden_size)` in the `best_model_structure.txt` file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "paths: /home/vmagent/app/e2eaiok/e2eAIOK/DeNas/asr/utils, /home/vmagent/app/e2eaiok/e2eAIOK/DeNas/asr\n",
      "['/home/vmagent/app/e2eaiok/e2eAIOK/DeNas', '/opt/intel/oneapi/advisor/2022.3.0/pythonapi', '/opt/intel/oneapi/intelpython/latest/envs/pytorch-1.12.0/lib/python39.zip', '/opt/intel/oneapi/intelpython/latest/envs/pytorch-1.12.0/lib/python3.9', '/opt/intel/oneapi/intelpython/latest/envs/pytorch-1.12.0/lib/python3.9/lib-dynload', '/opt/intel/oneapi/intelpython/latest/envs/pytorch-1.12.0/lib/python3.9/site-packages', '/opt/intel/oneapi/intelpython/latest/envs/pytorch-1.12.0/lib/python3.9/site-packages/e2eAIOK-0.2.1-py3.9.egg', '', '/home/vmagent/app/e2eaiok/e2eAIOK/DeNas', '/home/vmagent/app/e2eaiok/e2eAIOK/DeNas', '/home/vmagent/app/e2eaiok/e2eAIOK/DeNas', '/home/vmagent/app/e2eaiok/e2eAIOK/DeNas', '/home/vmagent/app/e2eaiok/e2eAIOK/DeNas', '/home/vmagent/app/e2eaiok/e2eAIOK/DeNas/asr']\n",
      "loading archive file /home/vmagent/app/dataset/bert-base-uncased\n",
      "12/01/2022 13:43:12 - INFO - nlp.supernet_bert -   Model config {\n",
      "  \"attention_probs_dropout_prob\": 0.1,\n",
      "  \"hidden_act\": \"gelu\",\n",
      "  \"hidden_dropout_prob\": 0.1,\n",
      "  \"hidden_size\": 768,\n",
      "  \"initializer_range\": 0.02,\n",
      "  \"intermediate_size\": 3072,\n",
      "  \"layer_norm_eps\": 1e-12,\n",
      "  \"max_position_embeddings\": 512,\n",
      "  \"num_attention_heads\": 12,\n",
      "  \"num_hidden_layers\": 12,\n",
      "  \"type_vocab_size\": 2,\n",
      "  \"vocab_size\": 30522\n",
      "}\n",
      "\n",
      "Weights of SuperBertModel not initialized from pretrained model: ['bert.dense_fit.weight', 'bert.dense_fit.bias']\n",
      "Weights from pretrained model not used in SuperBertModel: ['bert.cls.predictions.bias', 'bert.cls.predictions.transform.dense.weight', 'bert.cls.predictions.transform.dense.bias', 'bert.cls.predictions.decoder.weight', 'bert.cls.seq_relationship.weight', 'bert.cls.seq_relationship.bias', 'bert.cls.predictions.transform.LayerNorm.weight', 'bert.cls.predictions.transform.LayerNorm.bias']\n",
      "12/01/2022 13:43:14 - INFO - DENAS -   epoch = 0\n",
      "12/01/2022 13:43:18 - INFO - DENAS -   random 1/10 structure (11, 10, 640, 720, 992) nas_score 224.1220245361328 params 58.934512\n",
      "12/01/2022 13:43:21 - INFO - DENAS -   random 2/10 structure (6, 8, 512, 752, 2816) nas_score 173.40965270996094 params 58.612176\n",
      "12/01/2022 13:43:24 - INFO - DENAS -   random 3/10 structure (12, 11, 704, 656, 1248) nas_score 219.91477966308594 params 62.695536\n",
      "12/01/2022 13:43:28 - INFO - DENAS -   random 4/10 structure (10, 12, 768, 640, 1376) nas_score 254.247802734375 params 57.62336\n",
      "12/01/2022 13:43:31 - INFO - DENAS -   random 5/10 structure (8, 10, 640, 720, 2720) nas_score 142.49546813964844 params 69.01816\n",
      "12/01/2022 13:43:35 - INFO - DENAS -   random 6/10 structure (11, 11, 704, 704, 1824) nas_score 287.5343933105469 params 72.494048\n",
      "12/01/2022 13:43:40 - INFO - DENAS -   random 7/10 structure (11, 11, 704, 768, 1888) nas_score 230.45338439941406 params 80.21168\n",
      "12/01/2022 13:43:44 - INFO - DENAS -   random 8/10 structure (12, 10, 640, 576, 2112) nas_score 200.59535217285156 params 65.191104\n",
      "12/01/2022 13:43:48 - INFO - DENAS -   random 9/10 structure (10, 11, 704, 656, 2144) nas_score 248.85580444335938 params 67.47608\n",
      "12/01/2022 13:43:52 - INFO - DENAS -   random 10/10 structure (9, 11, 704, 624, 2720) nas_score 181.30479431152344 params 66.200592\n",
      "12/01/2022 13:43:52 - INFO - DENAS -   random_num = 10\n",
      "12/01/2022 13:43:57 - INFO - DENAS -   mutation 1/10 structure (12, 12, 768, 672, 1856) nas_score 233.03652954101562 params 76.114272\n",
      "12/01/2022 13:44:01 - INFO - DENAS -   mutation 2/10 structure (12, 10, 640, 576, 2528) nas_score 204.45196533203125 params 70.94688\n",
      "12/01/2022 13:44:04 - INFO - DENAS -   mutation 3/10 structure (11, 11, 704, 768, 864) nas_score 283.5840759277344 params 62.898912\n",
      "12/01/2022 13:44:09 - INFO - DENAS -   mutation 4/10 structure (12, 9, 576, 736, 2112) nas_score 224.6094512939453 params 81.140768\n",
      "12/01/2022 13:44:13 - INFO - DENAS -   mutation 5/10 structure (11, 11, 704, 736, 1824) nas_score 309.4450988769531 params 75.810816\n",
      "12/01/2022 13:44:16 - INFO - DENAS -   mutation 6/10 structure (12, 11, 704, 656, 864) nas_score 171.3025665283203 params 56.645232\n",
      "12/01/2022 13:44:23 - INFO - DENAS -   mutation 7/10 structure (12, 12, 768, 544, 2976) nas_score 174.1815948486328 params 76.192352\n",
      "12/01/2022 13:44:28 - INFO - DENAS -   mutation 8/10 structure (11, 11, 704, 768, 2688) nas_score 273.32733154296875 params 93.73728\n",
      "12/01/2022 13:44:32 - INFO - DENAS -   mutation 9/10 structure (11, 11, 704, 768, 1216) nas_score 303.4856872558594 params 68.850176\n",
      "12/01/2022 13:44:37 - INFO - DENAS -   mutation 10/10 structure (11, 10, 640, 720, 2464) nas_score 262.8741760253906 params 82.267184\n",
      "12/01/2022 13:44:37 - INFO - DENAS -   mutation_num = 10\n",
      "12/01/2022 13:44:43 - INFO - DENAS -   crossover 1/10 structure (11, 11, 704, 624, 2720) nas_score 248.56190490722656 params 76.521232\n",
      "12/01/2022 13:44:47 - INFO - DENAS -   crossover 2/10 structure (10, 11, 704, 656, 1824) nas_score 260.8335876464844 params 63.27448\n",
      "12/01/2022 13:44:51 - INFO - DENAS -   crossover 3/10 structure (10, 11, 704, 640, 2144) nas_score 241.08465576171875 params 65.82112\n",
      "12/01/2022 13:44:55 - INFO - DENAS -   crossover 4/10 structure (11, 11, 704, 624, 1888) nas_score 264.8471374511719 params 65.090384\n",
      "12/01/2022 13:44:58 - INFO - DENAS -   crossover 5/10 structure (6, 11, 704, 752, 2816) nas_score 169.1065216064453 params 62.080848\n",
      "12/01/2022 13:45:02 - INFO - DENAS -   crossover 6/10 structure (11, 11, 704, 768, 1824) nas_score 303.0057067871094 params 79.129632\n",
      "12/01/2022 13:45:06 - INFO - DENAS -   crossover 7/10 structure (10, 10, 640, 656, 2112) nas_score 245.7474365234375 params 65.37464\n",
      "12/01/2022 13:45:10 - INFO - DENAS -   crossover 8/10 structure (6, 12, 768, 752, 2816) nas_score 161.86341857910156 params 63.237072\n",
      "12/01/2022 13:45:15 - INFO - DENAS -   crossover 9/10 structure (8, 10, 640, 720, 2816) nas_score 136.3217315673828 params 70.124848\n",
      "12/01/2022 13:45:18 - INFO - DENAS -   crossover 10/10 structure (11, 11, 704, 640, 1824) nas_score 275.316650390625 params 65.866656\n",
      "12/01/2022 13:45:18 - INFO - DENAS -   crossover_num = 10\n",
      "DE-NAS search best structure took 124.73312006797642 sec\n",
      "12/01/2022 13:45:18 - INFO - DENAS -   best structure (11, 11, 704, 736, 1824) nas_score 309.4450988769531 params 75.810816\n",
      "DE-NAS completed, best structure is (11, 11, 704, 736, 1824)\n"
     ]
    }
   ],
   "source": [
    "from e2eAIOK.DeNas.search.utils import parse_config\n",
    "from e2eAIOK.DeNas.thirdparty.supernet_hf import SuperHFModel\n",
    "from e2eAIOK.DeNas.search.SearchEngineFactory import SearchEngineFactory\n",
    "\n",
    "\n",
    "# parse DE-NAS search configure\n",
    "params = parse_config('/home/vmagent/app/e2eaiok/conf/denas/hf/e2eaiok_denas_hf.conf')\n",
    "\n",
    "# construct supernet and search space\n",
    "super_net = SuperHFModel.from_pretrained(params.supernet)\n",
    "search_space = SuperHFModel.search_space_generation(params.supernet)\n",
    "\n",
    "# create DE-NAS searcher\n",
    "searcher = SearchEngineFactory.create_search_engine(params = params, super_net = super_net, search_space = search_space)\n",
    "\n",
    "# trigger the search process\n",
    "searcher.search()\n",
    "best_structure = searcher.get_best_structures()\n",
    "print(f\"DE-NAS completed, best structure is {best_structure}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Launch Training with Best Searched Model Structure"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Launch the DE-NAS training process with user-self defined training process. Below is an simple instroduction how to add DE-NAS constructed model into the customer's training pipeline or DE-NAS training pipeline (etc., fine-tuning DE-NAS HF BERT in the SQuADv1.1 task).\n",
    "\n",
    "Noted: Below performance result is using sample dataset and small iterations to demonstrate its function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'domain': 'hf', 'task_name': 'squad1', 'task_type': 'classification', 'supernet': 'bert-base-uncased', 'tokenizer': 'bert-base-uncased', 'optimizer': 'BertAdam', 'criterion': 'CrossEntropyQALoss', 'lr_scheduler': 'warmup_linear', 'eval_metric': 'qa_f1', 'dist_backend': 'gloo', 'input_id': 'input_ids attention_mask token_type_ids', 'data_set': 'SQuADv1.1', 'best_model_structure': '/home/vmagent/app/e2eaiok/e2eAIOK/DeNas/best_model_structure.txt', 'model': '/home/vmagent/app/dataset/bert-base-uncased/', 'model_dir': '/home/vmagent/app/dataset/bert-base-uncased/', 'data_dir': '/home/vmagent/app/dataset/SQuAD_small/', 'output_dir': '/home/vmagent/app/e2eaiok/e2eAIOK/DeNas/thirdparty/', 'hidden_size': 640, 'gradient_accumulation_steps': 1, 'warmup_proportion': 0.1, 'learning_rate': 6e-05, 'weight_decay': 0.01, 'initializer_range': 0.02, 'train_epochs': 2, 'max_seq_length': 384, 'doc_stride': 128, 'train_batch_size': 32, 'eval_batch_size': 8, 'eval_step': 500, 'n_best_size': 20, 'max_answer_length': 30, 'max_query_length': 64, 'version_2_with_negative': 0, 'null_score_diff_threshold': 0.0, 'num_labels': 2, 'num_workers': 1, 'pin_mem': True, 'verbose_logging': False, 'no_cuda': True, 'do_lower_case': True, 'metric_threshold': 'None'}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at /home/vmagent/app/dataset/bert-base-uncased/ were not used when initializing BertModel: ['cls.predictions.bias', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.seq_relationship.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.seq_relationship.weight', 'cls.predictions.decoder.weight']\n",
      "- This IS expected if you are initializing BertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "03/18/2023 11:17:56 - INFO - e2eAIOK.DeNas.module.nlp.tokenization -   loading vocabulary file\n",
      "03/18/2023 11:17:56 - INFO - e2eAIOK.common.trainer.data.nlp.data_builder_squad -   load 1027 examples!\n",
      "03/18/2023 11:17:58 - INFO - e2eAIOK.DeNas.module.nlp.tokenization -   loading vocabulary file\n",
      "03/18/2023 11:17:58 - INFO - e2eAIOK.common.trainer.data.nlp.data_builder_squad -   load 1680 examples!\n",
      "03/18/2023 11:18:01 - INFO - Trainer -   Trainer config: {'domain': 'hf', 'task_name': 'squad1', 'task_type': 'classification', 'supernet': 'bert-base-uncased', 'tokenizer': 'bert-base-uncased', 'optimizer': 'BertAdam', 'criterion': 'CrossEntropyQALoss', 'lr_scheduler': 'warmup_linear', 'eval_metric': 'qa_f1', 'dist_backend': 'gloo', 'input_id': 'input_ids attention_mask token_type_ids', 'data_set': 'SQuADv1.1', 'best_model_structure': '/home/vmagent/app/e2eaiok/e2eAIOK/DeNas/best_model_structure.txt', 'model': '/home/vmagent/app/dataset/bert-base-uncased/', 'model_dir': '/home/vmagent/app/dataset/bert-base-uncased/', 'data_dir': '/home/vmagent/app/dataset/SQuAD_small/', 'output_dir': '/home/vmagent/app/e2eaiok/e2eAIOK/DeNas/thirdparty/', 'hidden_size': 640, 'gradient_accumulation_steps': 1, 'warmup_proportion': 0.1, 'learning_rate': 6e-05, 'weight_decay': 0.01, 'initializer_range': 0.02, 'train_epochs': 2, 'max_seq_length': 384, 'doc_stride': 128, 'train_batch_size': 32, 'eval_batch_size': 8, 'eval_step': 500, 'n_best_size': 20, 'max_answer_length': 30, 'max_query_length': 64, 'version_2_with_negative': 0, 'null_score_diff_threshold': 0.0, 'num_labels': 2, 'num_workers': 1, 'pin_mem': True, 'verbose_logging': False, 'no_cuda': True, 'do_lower_case': True, 'metric_threshold': 'None', 'num_train_steps': 34}\n",
      "03/18/2023 11:18:01 - INFO - root -   Please set the 'teacher_model' if you want to use the KD within the BERT training process\n",
      "Iteration:   0%|                                                                                                                       | 0/34 [00:00<?, ?it/s]/home/vmagent/app/e2eaiok/e2eAIOK/DeNas/module/nlp/optimization.py:249: UserWarning: This overload of add_ is deprecated:\n",
      "\tadd_(Number alpha, Tensor other)\n",
      "Consider using one of the following signatures instead:\n",
      "\tadd_(Tensor other, *, Number alpha) (Triggered internally at  ../torch/csrc/utils/python_arg_parser.cpp:1174.)\n",
      "  next_m.mul_(beta1).add_(1 - beta1, grad)\n",
      "Iteration: 100%|##############################################################################################################| 34/34 [01:56<00:00,  3.43s/it]\n",
      "03/18/2023 11:19:58 - INFO - Trainer -   Epoch 0 training time:116.60106873512268\n",
      "Iteration:  97%|##########################################################################################################7   | 33/34 [01:43<00:03,  3.12s/it]03/18/2023 11:21:42 - INFO - Trainer -   ***** Running evaluation *****\n",
      "03/18/2023 11:21:42 - INFO - Trainer -     Epoch = 1 iter 68 step\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "***** Eval results *****\n",
      "em = 7.559523809523809\n",
      "infer_cnt = 213\n",
      "infer_time = 140.6275164319249\n",
      "qa_f1 = 13.243573213106428\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|##############################################################################################################| 34/34 [02:21<00:00,  4.15s/it]\n",
      "03/18/2023 11:22:19 - INFO - Trainer -   Epoch 1 training time:141.16966462135315\n",
      "03/18/2023 11:22:19 - INFO - Trainer -   **************S*************\n",
      "task_name = squad1\n",
      "total training time = 257.77339577674866\n",
      "best_acc = 13.243573213106428\n",
      "**************E*************\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from e2eAIOK.DeNas.thirdparty.supernet_hf import SuperHFModel\n",
    "from e2eAIOK.DeNas.thirdparty.utils import decode_arch\n",
    "from e2eAIOK.DeNas.search.utils import parse_config\n",
    "from e2eAIOK.common.trainer.data.nlp.data_builder_squad import DataBuilderSQuAD\n",
    "from e2eAIOK.DeNas.nlp.utils import bert_create_optimizer, bert_create_criterion, bert_create_scheduler, bert_create_metric\n",
    "from e2eAIOK.DeNas.nlp.bert_trainer import BERTTrainer\n",
    "\n",
    "from easydict import EasyDict as edict\n",
    "from torch import nn\n",
    "\n",
    "### Adding a task-specific layer to the last layer of denas model ###\n",
    "class SuperHFTaskModel(nn.Module):\n",
    "    def __init__(self, encoder_model, cfg):\n",
    "        super(SuperHFTaskModel, self).__init__()\n",
    "        self.cfg = cfg\n",
    "        self.encoder_model = encoder_model\n",
    "        if self.cfg[\"task_type\"] == \"classification\":\n",
    "            self.output = nn.Linear(self.cfg[\"hidden_size\"], self.cfg[\"num_labels\"])\n",
    "        else:\n",
    "            raise NotImplementedError(\"Task Type are not supported yet!\")\n",
    "        self.output.apply(self.init_output_weight)\n",
    "    \n",
    "    def init_output_weight(self, module):\n",
    "        if isinstance(module, nn.Linear):\n",
    "            module.weight.data.normal_(mean=0.0, std=self.cfg.initializer_range)\n",
    "            if hasattr(module, \"bias\") and module.bias is not None:\n",
    "                module.bias.data.zero_()\n",
    "\n",
    "    def forward(self, x):\n",
    "        input_keys = self.cfg[\"input_id\"].strip().split()\n",
    "        item = x.split(1, -1)\n",
    "        inputs = dict()\n",
    "        for id, input_key in enumerate(input_keys):\n",
    "            inputs[input_key] = item[id].squeeze(-1)\n",
    "        output = self.encoder_model(**inputs)\n",
    "        last_hidden_state = output.last_hidden_state\n",
    "        logits = self.output(last_hidden_state)\n",
    "        return logits\n",
    "\n",
    "### Model Contruction with DE-NAS Configuration ###\n",
    "model_arch = decode_arch(\"/home/vmagent/app/e2eaiok/e2eAIOK/DeNas/best_model_structure.txt\")\n",
    "cfg = edict(parse_config(\"/home/vmagent/app/e2eaiok/conf/denas/hf/e2eaiok_denas_train_bert.conf\"))\n",
    "print(cfg)\n",
    "model = SuperHFModel.set_sample_config(\"/home/vmagent/app/dataset/bert-base-uncased/\", **model_arch)\n",
    "model = SuperHFTaskModel(model, cfg)\n",
    "\n",
    "### Training Pipeline Construction ###\n",
    "train_dataloader, eval_dataloader, other_data = DataBuilderSQuAD(cfg).get_dataloader()\n",
    "cfg.num_train_steps = len(train_dataloader)\n",
    "optimizer = bert_create_optimizer(model, cfg)\n",
    "criterion = bert_create_criterion(cfg)\n",
    "scheduler = bert_create_scheduler(cfg)\n",
    "metric = bert_create_metric(cfg)\n",
    "\n",
    "### Create DE-NAS Trainer ###\n",
    "trainer = BERTTrainer(cfg, model, train_dataloader, eval_dataloader, other_data, optimizer, criterion, scheduler, metric)\n",
    "\n",
    "### Trigger the Training Process ###\n",
    "trainer.fit()"
   ]
  }
 ],
 "metadata": {
  "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.9.7"
  },
  "vscode": {
   "interpreter": {
    "hash": "b38dbbca6c1c1b09874e6226521bb418530c95443a0afe85f48e402a3f241a05"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}