full_pipeline.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import tensorflow as tf\n",
    "import tensorflow.keras as keras\n",
    "from utils import set_seed\n",
    "from plot_utils import plot_eval_results\n",
    "from converters import Converter\n",
    "from config import cfg\n",
    "from c_utils import write_model_h\n",
    "from utils import save_test_data\n",
    "import plot_utils as pu\n",
    "from model import save_pruned_model\n",
    "\n",
    "set_seed()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nyuv2_torch_ds_adapter import get_tf_nyuv2_ds\n",
    "import argparse\n",
    "args = argparse.Namespace()\n",
    "args.truncate_testset = False\n",
    "args.crop_size = (640, 480)\n",
    "args.target_size = cfg.img_size\n",
    "args.out_fold_ratio = 1\n",
    "args.is_maxim = False\n",
    "cfg.do_overfit=False\n",
    "args.batch_size=cfg.batch_size*4\n",
    "\n",
    "ds_train, ds_val, ds_test = get_tf_nyuv2_ds(cfg.base_dataset_dir, args)\n",
    "\n",
    "x_val= next(iter(ds_val))\n",
    "x_train= next(iter(ds_train))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_train[0].numpy().max(), x_train[0].numpy().min(), x_train[1].numpy().max(), x_train[1].numpy().min(), "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1, 2, figsize=(10, 5))\n",
    "axs[0].imshow(x_train[0][0].numpy())\n",
    "axs[1].imshow(x_train[1][0].numpy())\n",
    "axs[0].set_title('Image')\n",
    "axs[1].set_title('Depth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "converter = Converter(cfg)\n",
    "\n",
    "os.makedirs(cfg.save_model_dir, exist_ok=True)\n",
    "os.makedirs(cfg.save_cfiles_dir, exist_ok=True)\n",
    "os.makedirs(cfg.save_test_data_dir, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_names = [\n",
    "    \"depth_model_quant8_dynR\",\n",
    "    \"depth_full_quant\",\n",
    "    \"depth_qat_int8\",\n",
    "    \"pruned_model\",\n",
    "    \"pruned_model_unstructured\",\n",
    "    \"pruned_model_unstructured_dynamic\",\n",
    "    \"pruned_qat_model\",\n",
    "    \"depth_model_fp32\",\n",
    "]\n",
    "cfiles = {\n",
    "    \"depth_model_quant8_dynR\": \"depth_model_quant8_dynR\",\n",
    "    \"depth_full_quant\": \"q8depth\",\n",
    "    \"depth_qat_int8\": \"qat8depth\",\n",
    "    \"pruned_model\": \"pruned\",\n",
    "    \"pruned_model_unstructured\": \"pruned_unstructured\",\n",
    "    \"pruned_model_unstructured_dynamic\": \"pruned_unstructured_dynamic\",\n",
    "    \"pruned_qat_model\": \"pruned_unstructured_qat_model\",\n",
    "    \"depth_model_fp32\": \"depth_model_fp32\",\n",
    "}\n",
    "save_test_data(cfg.save_test_data_dir, x_train[0], x_train[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "notebookRunGroups": {
     "groupValue": "1"
    }
   },
   "outputs": [],
   "source": [
    "import pickle\n",
    "from eval import eval_model\n",
    "from loss import calculate_loss\n",
    "from model import get_model\n",
    "import tensorflow_model_optimization as tfmot\n",
    "\n",
    "\n",
    "def fit_eval(model, model_name, do_save_model=True, epochs=cfg.epochs, do_savefig=True):\n",
    "    keras.backend.clear_session()\n",
    "    metrics = tf.keras.metrics.Mean(name=\"loss\")\n",
    "\n",
    "    def custom_metric(y_true, y_pred, sample_weight=None):\n",
    "        metric_value = calculate_loss(y_true, y_pred)\n",
    "        metrics.update_state(metric_value, sample_weight=sample_weight)\n",
    "        return metric_value\n",
    "\n",
    "    optimizer = tf.keras.optimizers.Adam(learning_rate=0.0005)\n",
    "\n",
    "    model.compile(optimizer=optimizer, loss=\"mae\", metrics=[custom_metric])\n",
    "\n",
    "    reduce_lr = tf.keras.callbacks.ReduceLROnPlateau(\n",
    "        monitor=\"loss\", factor=0.2, min_lr=1e-7, patience=5, min_delta=1e-2, verbose=1\n",
    "    )\n",
    "    callbacks = [reduce_lr]\n",
    "    if \"pruned\" in model_name:\n",
    "        callbacks.append(tfmot.sparsity.keras.UpdatePruningStep())\n",
    "    if not cfg.do_overfit:\n",
    "        tbCallBack = tf.keras.callbacks.TensorBoard(\n",
    "            log_dir=f\"{cfg.logdir}/tb_logs\",\n",
    "            histogram_freq=0,\n",
    "            write_graph=False,\n",
    "            write_images=False,\n",
    "        )\n",
    "        es = tf.keras.callbacks.EarlyStopping(\n",
    "            patience=cfg.es_patience, \n",
    "            min_delta=5e-3, \n",
    "            monitor=\"loss\",\n",
    "        )\n",
    "        model_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(\n",
    "            filepath=cfg.logdir + f\"/{model_name}\",\n",
    "            save_weights_only=True,\n",
    "            monitor=\"custom_metric\",\n",
    "            mode=\"min\",\n",
    "        )\n",
    "        callbacks.append(es)\n",
    "        callbacks.append(model_checkpoint_callback)\n",
    "        callbacks.append(tbCallBack)\n",
    "    history = model.fit(\n",
    "        x=ds_train,\n",
    "        epochs=epochs,\n",
    "        validation_data=ds_val,\n",
    "        callbacks=callbacks,\n",
    "        verbose=1,\n",
    "    )\n",
    "    with open(f\"{cfg.save_model_dir}/{model_name}_history.pkl\", \"wb\") as f:\n",
    "        pickle.dump(history.history, f)\n",
    "\n",
    "    pu.plot_history(history, f\"{cfg.save_model_dir}/{model_name}_history.png\")\n",
    "    if not do_savefig:\n",
    "        plt.show()\n",
    "    tflite_path=None\n",
    "    model_tflite=None\n",
    "    if do_save_model:\n",
    "        model.save(f\"{cfg.save_model_dir}/{model_name}.h5\")\n",
    "        model_tflite, tflite_path = converter.keras_to_tflite(model, model_name, do_return_path=True)\n",
    "    metrics = eval_model(\n",
    "        model=model,\n",
    "        test_ds=ds_val,\n",
    "        tflite_path=tflite_path,\n",
    "        model_name=model_name,\n",
    "        metrics_file_path=f\"{cfg.save_model_dir}/metrics.json\",\n",
    "    )\n",
    "    return model_tflite, metrics, history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build model\n",
    "keras.backend.clear_session()\n",
    "fp_model = get_model(\n",
    "    cfg.img_size, cfg.num_classes, in_channels=cfg.in_channels, use_qat=False,\n",
    "    do_downsample_model=True\n",
    ")\n",
    "fit_eval(fp_model, model_names[7], epochs=cfg.epochs)\n",
    "model_tflite, tflite_path = converter.keras_to_tflite(fp_model, model_names[7], do_return_path=True)\n",
    "write_model_h(cfiles[model_names[7]], model_tflite, cfg.save_cfiles_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "out = fp_model.predict(x_train[0][0][None, ...])\n",
    "plot_eval_results(out[0], x_train[1][0], x_train[0][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dynR_quant_tflite_model = converter.dynamic_range_quantization(fp_model, model_names[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tflite_model_quant_int8=converter.eight_bit_quantization(fp_model, ds_train, model_name=model_names[1])\n",
    "converter.check_quantized_model(tflite_model_quant_int8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "depth_full_quant_tflite_path=f'{cfg.save_model_dir}/{model_names[1]}.tflite'\n",
    "eval_model(\n",
    "        ds_val,\n",
    "        tflite_path=depth_full_quant_tflite_path,\n",
    "        model=None,\n",
    "        model_name=model_names[1],\n",
    "        metrics_file_path=f\"{cfg.save_model_dir}/metrics.json\",\n",
    "    )\n",
    "write_model_h(cfiles[model_names[1]], tflite_model_quant_int8, cfg.save_cfiles_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert the model to a quantization aware model\n",
    "quant_aware_model = get_model(\n",
    "    cfg.img_size, cfg.num_classes, in_channels=cfg.in_channels, use_qat=True\n",
    ")\n",
    "model_name = model_names[2]\n",
    "fit_eval(quant_aware_model, model_name)\n",
    "quant_aware_model_converted = converter.eight_bit_quantization(\n",
    "    quant_aware_model, ds_train, model_name=model_name\n",
    ")\n",
    "write_model_h(cfiles[model_name], quant_aware_model_converted, cfg.save_cfiles_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = model_names[3]\n",
    "pruned_model = get_model(cfg.img_size, cfg.num_classes, in_channels=cfg.in_channels, use_qat=False, use_pruning=True,\n",
    "                         use_pruning_struct=True)\n",
    "fit_eval(pruned_model, model_name, do_save_model=False)\n",
    "pruned_model_for_export = save_pruned_model(pruned_model, f\"{cfg.save_model_dir}/{model_name}.h5\")\n",
    "pruned_tflite_model = converter.keras_to_tflite(pruned_model_for_export, model_name)\n",
    "write_model_h(cfiles[model_name], pruned_tflite_model, cfg.save_cfiles_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = model_names[4]\n",
    "pruned_model_unstruct = get_model(\n",
    "    cfg.img_size,\n",
    "    cfg.num_classes,\n",
    "    in_channels=cfg.in_channels,\n",
    "    use_qat=False,\n",
    "    use_pruning=True,\n",
    "    use_pruning_struct=False,\n",
    "    use_dynamic_sparsity=False,\n",
    ")\n",
    "fit_eval(pruned_model_unstruct, model_name, do_save_model=False)\n",
    "pruned_model_for_export = save_pruned_model(\n",
    "    pruned_model_unstruct, f\"{cfg.save_model_dir}/{model_name}.h5\"\n",
    ")\n",
    "pruned_tflite_model = converter.keras_to_tflite(pruned_model_for_export, model_name)\n",
    "eval_model(\n",
    "        model=pruned_model_unstruct,\n",
    "        test_ds=ds_val,\n",
    "        tflite_path=f\"{cfg.save_model_dir}/{model_name}.tflite\",\n",
    "        model_name=model_name,\n",
    "        metrics_file_path=f\"{cfg.save_model_dir}/metrics.json\",\n",
    "    )\n",
    "write_model_h(cfiles[model_name], pruned_tflite_model, cfg.save_cfiles_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = model_names[5]\n",
    "pruned_model_unstruct_dynamic = get_model(\n",
    "    cfg.img_size,\n",
    "    cfg.num_classes,\n",
    "    in_channels=cfg.in_channels,\n",
    "    use_qat=False,\n",
    "    use_pruning=True,\n",
    "    use_pruning_struct=False,\n",
    "    use_dynamic_sparsity=True,\n",
    ")\n",
    "fit_eval(pruned_model_unstruct_dynamic, model_name, do_save_model=False)\n",
    "pruned_model_unstructured_for_export = save_pruned_model(\n",
    "    pruned_model_unstruct_dynamic, f\"{cfg.save_model_dir}/{model_name}.h5\"\n",
    ")\n",
    "pruned_tflite_model = converter.keras_to_tflite(pruned_model_unstructured_for_export, model_name)\n",
    "eval_model(\n",
    "        model=pruned_model_unstruct_dynamic,\n",
    "        test_ds=ds_val,\n",
    "        tflite_path=f\"{cfg.save_model_dir}/{model_name}.tflite\",\n",
    "        model_name=model_name,\n",
    "        metrics_file_path=f\"{cfg.save_model_dir}/metrics.json\",\n",
    "    )\n",
    "write_model_h(cfiles[model_name], pruned_tflite_model, cfg.save_cfiles_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = model_names[6]\n",
    "pruned_qat_model = get_model(\n",
    "    cfg.img_size,\n",
    "    cfg.num_classes,\n",
    "    in_channels=cfg.in_channels,\n",
    "    use_qat=True,\n",
    "    use_pruning=True,\n",
    "    use_pruning_struct=False,\n",
    "    use_dynamic_sparsity=False,\n",
    "    pruned_model_unstructured_for_export=pruned_model_unstructured_for_export\n",
    ")\n",
    "fit_eval(pruned_qat_model, model_name, do_save_model=False)\n",
    "pruned_model_for_export = save_pruned_model(\n",
    "    pruned_qat_model, f\"{cfg.save_model_dir}/{model_name}.h5\"\n",
    ")\n",
    "pruned_tflite_model = converter.keras_to_tflite(pruned_model_for_export, model_name)\n",
    "eval_model(\n",
    "        model=pruned_qat_model,\n",
    "        test_ds=ds_val,\n",
    "        tflite_path=f\"{cfg.save_model_dir}/{model_name}.tflite\",\n",
    "        model_name=model_name,\n",
    "        metrics_file_path=f\"{cfg.save_model_dir}/metrics.json\",\n",
    "    )\n",
    "write_model_h(cfiles[model_name], pruned_tflite_model, cfg.save_cfiles_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sc",
   "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.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}