Training with MXNet DataLoader tf.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TensorFlow ArcFace"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Ref https://github.com/peteryuX/arcface-tf2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.6/dist-packages/torch/cuda/__init__.py:104: UserWarning: \n",
      "A100-SXM4-40GB with CUDA capability sm_80 is not compatible with the current PyTorch installation.\n",
      "The current PyTorch install supports CUDA capabilities sm_37 sm_50 sm_60 sm_70 sm_75.\n",
      "If you want to use the A100-SXM4-40GB GPU with PyTorch, please check the instructions at https://pytorch.org/get-started/locally/\n",
      "\n",
      "  warnings.warn(incompatible_device_warn.format(device_name, capability, \" \".join(arch_list), device_name))\n"
     ]
    }
   ],
   "source": [
    "from insightface.recognition.arcface_torch.dataset import MXFaceDataset, DataLoaderX\n",
    "import torch.utils.data.distributed\n",
    "import torch.distributed as dist\n",
    "from easydict import EasyDict as edict\n",
    "import time\n",
    "import os\n",
    "\n",
    "cfg = edict()\n",
    "cfg.rec = \"/raid/workspace/honghee/faces_emore/\"\n",
    "local_rank = 0\n",
    "cfg.batch_size = 64\n",
    "\n",
    "try:\n",
    "    world_size = int(os.environ['WORLD_SIZE'])\n",
    "    rank = int(os.environ['RANK'])\n",
    "    dist_url = \"tcp://{}:{}\".format(os.environ[\"MASTER_ADDR\"], os.environ[\"MASTER_PORT\"])\n",
    "except KeyError:\n",
    "    world_size = 1\n",
    "    rank = 0\n",
    "    dist_url = \"tcp://127.0.0.1:12585\"\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"2\"\n",
    "dist.init_process_group(backend='nccl', init_method=dist_url, rank=rank, world_size=world_size)\n",
    "torch.cuda.set_device(local_rank)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2\r\n"
     ]
    }
   ],
   "source": [
    "!echo $CUDA_VISIBLE_DEVICES"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_set = MXFaceDataset(root_dir=cfg.rec, local_rank=local_rank)\n",
    "train_sampler = torch.utils.data.distributed.DistributedSampler(\n",
    "    train_set, shuffle=True)\n",
    "train_loader = DataLoaderX(\n",
    "    local_rank=local_rank, dataset=train_set, batch_size=cfg.batch_size,\n",
    "    sampler=train_sampler, num_workers=2, pin_memory=True, drop_last=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for step, (x_batch_train, y_batch_train) in enumerate(train_loader):\n",
    "#     x_batch_train = x_batch_train.cpu().numpy()\n",
    "#     y_batch_train = y_batch_train.cpu().numpy()\n",
    "#     print(y_batch_train)\n",
    "#     time.sleep(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.callbacks import LearningRateScheduler, EarlyStopping, ModelCheckpoint, ReduceLROnPlateau, TensorBoard\n",
    "from models import ArcFaceModel\n",
    "# from losses import SoftmaxLoss\n",
    "from losses import softmax_loss\n",
    "import dataset\n",
    "import tensorflow as tf\n",
    "import os\n",
    "import logging\n",
    "\n",
    "tf.get_logger().setLevel(logging.ERROR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CurHead\n",
      "EfficientNetB6\n",
      "projection head: False\n",
      "Adam: False\n",
      "epoch: 100\n",
      "batch size: 64\n"
     ]
    }
   ],
   "source": [
    "### IJB-C Dataset\n",
    "# batch_size = 128\n",
    "# input_size = 112\n",
    "# embd_shape = 512\n",
    "# head_type = 'ArcHead'\n",
    "# backbone_type = 'MobileNetV2'\n",
    "# w_decay=5e-4\n",
    "# num_classes = 3584 \n",
    "# base_lr = 0.01\n",
    "# dataset_len = 13033 \n",
    "# epochs = 100\n",
    "# steps_per_epoch = dataset_len // batch_size\n",
    "\n",
    "### MS1M dataset\n",
    "num_classes = 85742 \n",
    "dataset_len = 5822653\n",
    "batch_size = 64 # Initially 128\n",
    "input_size = 112\n",
    "embd_shape = 512\n",
    "train_size = int(0.8 * dataset_len)\n",
    "steps_per_epoch = train_size // batch_size\n",
    "val_size = dataset_len - train_size\n",
    "validation_steps = val_size // batch_size\n",
    "\n",
    "w_decay=5e-4\n",
    "epochs = 100\n",
    "\n",
    "save_steps = 1000\n",
    "steps = 1\n",
    "is_ccrop=False\n",
    "binary_img=True\n",
    "\n",
    "is_Adam = False   # True\n",
    "projection_head = False  # True\n",
    "dgx = True\n",
    "\n",
    "head_type = 'CurHead' # ''ArcHead', CosHead', 'SphereHead', 'CurHead', 'AdaHead', CadHead'  \n",
    "# Backbones w/ pretrained weights:\n",
    "#     MobileNet, MobileNetV2, InceptionResNetV2, InceptionV3, ResNet50, ResNet50V2, ResNet101V2, NASNetLarge, NASNetMobile, Xception, MobileNetV3Large, MobileNetV3Small, EfficientNetLite0~6, EfficientNetB0~7\n",
    "# Backbones w/o pretrained weights:\n",
    "#      MnasNetA1, MnasNetB1, MnasNetSmall \n",
    "backbone_type = 'EfficientNetB6' \n",
    "\n",
    "if head_type == 'SphereHead':\n",
    "    base_lr = 0.01 \n",
    "    margin = 1.35\n",
    "    logist_scale = 30.0 \n",
    "elif head_type == 'CosHead':\n",
    "    base_lr = 0.01 \n",
    "    margin=0.35\n",
    "    logist_scale=64\n",
    "elif head_type == 'ArcHead':\n",
    "    base_lr = 0.1 \n",
    "    margin=0.5\n",
    "    logist_scale=64\n",
    "elif head_type == 'CurHead': \n",
    "    base_lr = 0.1 \n",
    "    margin=0.5\n",
    "    logist_scale=64\n",
    "elif head_type == 'CadHead': \n",
    "    base_lr = 0.1 \n",
    "    margin=0.0\n",
    "    logist_scale=64\n",
    "elif head_type == 'AdaHead':\n",
    "    base_lr = 0.1 \n",
    "    margin=0.0\n",
    "    logist_scale=64\n",
    "else:\n",
    "    base_lr = 0.01 # initially 0.01\n",
    "    \n",
    "print(head_type)\n",
    "print(backbone_type)\n",
    "print(\"projection head:\", projection_head)\n",
    "print(\"Adam:\", is_Adam)\n",
    "print(\"epoch:\", epochs)\n",
    "print(\"batch size:\", batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"arcface_model\"\n",
      "__________________________________________________________________________________________________\n",
      "Layer (type)                    Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      "input_image (InputLayer)        [(None, 112, 112, 3) 0                                            \n",
      "__________________________________________________________________________________________________\n",
      "efficientnetb6 (Functional)     (None, 4, 4, 2304)   40960136    input_image[0][0]                \n",
      "__________________________________________________________________________________________________\n",
      "OutputLayer (Functional)        (None, 512)          18886144    efficientnetb6[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "label (InputLayer)              [(None,)]            0                                            \n",
      "__________________________________________________________________________________________________\n",
      "CurHead (Functional)            (None, 85742)        43899905    OutputLayer[0][0]                \n",
      "                                                                 label[0][0]                      \n",
      "==================================================================================================\n",
      "Total params: 103,746,185\n",
      "Trainable params: 103,516,120\n",
      "Non-trainable params: 230,065\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "model = ArcFaceModel(size=input_size,\n",
    "                         backbone_type=backbone_type,\n",
    "                         num_classes=num_classes,\n",
    "                         margin=margin, \n",
    "                         logist_scale=logist_scale,\n",
    "                         head_type=head_type,\n",
    "                         embd_shape=embd_shape,\n",
    "                         w_decay=w_decay,\n",
    "                         training=True,\n",
    "                         projection_head=projection_head)\n",
    "model.summary()\n",
    "\n",
    "learning_rate = tf.constant(base_lr)\n",
    "lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(learning_rate,\n",
    "                                                             decay_steps=300000, \n",
    "                                               decay_rate=0.1,staircase=True)\n",
    "if is_Adam:\n",
    "    optimizer = tf.keras.optimizers.Adam(\n",
    "        learning_rate=lr_schedule)\n",
    "else:\n",
    "#         optimizer = tf.keras.optimizers.SGD(\n",
    "#             learning_rate=learning_rate, momentum=0.9)\n",
    "    optimizer = tf.keras.optimizers.SGD(\n",
    "        learning_rate=lr_schedule, momentum=0.9)\n",
    "\n",
    "# checkpoint = tf.train.Checkpoint(optimizer=optimizer, model=model)\n",
    "# manager = tf.train.CheckpointManager(\n",
    "#     checkpoint, directory=\"tmp/ckpt\", max_to_keep=5)\n",
    "# status = checkpoint.restore(manager.latest_checkpoint)\n",
    "\n",
    "# Prepare the metrics.\n",
    "train_acc_metric = tf.keras.metrics.SparseCategoricalAccuracy()\n",
    "val_acc_metric = tf.keras.metrics.SparseCategoricalAccuracy()\n",
    "\n",
    "#     model.compile(optimizer=optimizer, loss=loss_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "======================================================\n",
      "================= Start of epoch 0 ==================\n",
      "======================================================\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "Input 0 is incompatible with layer arcface_model: expected shape=(None, 112, 112, 3), found shape=(64, 3, 112, 112)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-8-ce8e3ce2ed18>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     32\u001b[0m         \u001b[0my_batch_train\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0my_batch_train\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcpu\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnumpy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     33\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0mtf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mGradientTape\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mtape\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 34\u001b[0;31m             \u001b[0mlogits\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx_batch_train\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my_batch_train\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtraining\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     35\u001b[0m             \u001b[0mloss_value\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mloss_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0my_batch_train\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlogits\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     36\u001b[0m         \u001b[0mgrads\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtape\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgradient\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mloss_value\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrainable_weights\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/base_layer.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    996\u001b[0m         \u001b[0minputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_maybe_cast_inputs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput_list\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    997\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 998\u001b[0;31m       \u001b[0minput_spec\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0massert_input_compatibility\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minput_spec\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    999\u001b[0m       \u001b[0;32mif\u001b[0m \u001b[0meager\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1000\u001b[0m         \u001b[0mcall_fn\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcall\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/input_spec.py\u001b[0m in \u001b[0;36massert_input_compatibility\u001b[0;34m(input_spec, inputs, layer_name)\u001b[0m\n\u001b[1;32m    272\u001b[0m                              \u001b[0;34m' is incompatible with layer '\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mlayer_name\u001b[0m \u001b[0;34m+\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    273\u001b[0m                              \u001b[0;34m': expected shape='\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mspec\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mshape\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m+\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 274\u001b[0;31m                              ', found shape=' + display_shape(x.shape))\n\u001b[0m\u001b[1;32m    275\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    276\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mValueError\u001b[0m: Input 0 is incompatible with layer arcface_model: expected shape=(None, 112, 112, 3), found shape=(64, 3, 112, 112)"
     ]
    }
   ],
   "source": [
    "import time\n",
    "from absl import app, flags, logging\n",
    "from absl.flags import FLAGS\n",
    "import cv2\n",
    "import os\n",
    "\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from glob import glob\n",
    "\n",
    "from evaluations import get_val_data, perform_val\n",
    "from models import ArcFaceModel\n",
    "from utils import set_memory_growth, load_yaml\n",
    "\n",
    "\n",
    "lfw, agedb_30, cfp_fp, lfw_issame, agedb_30_issame, cfp_fp_issame = get_val_data(\"/raid/workspace/honghee/\")\n",
    "\n",
    "loss_fn = softmax_loss\n",
    "epochs = 25\n",
    "\n",
    "best_val_acc = 0\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    print(\"======================================================\")\n",
    "    print(\"================= Start of epoch %d ==================\" % (epoch,))\n",
    "    print(\"======================================================\")\n",
    "    start_time = time.time()\n",
    "\n",
    "    # Iterate over the batches of the dataset.\n",
    "    for step, (x_batch_train, y_batch_train) in enumerate(train_loader):\n",
    "        x_batch_train = x_batch_train.cpu().numpy()\n",
    "        y_batch_train = y_batch_train.cpu().numpy()\n",
    "        with tf.GradientTape() as tape:\n",
    "            logits = model((x_batch_train, y_batch_train), training=True)\n",
    "            loss_value = loss_fn(y_batch_train, logits)\n",
    "        grads = tape.gradient(loss_value, model.trainable_weights)\n",
    "        optimizer.apply_gradients(zip(grads, model.trainable_weights))\n",
    "\n",
    "        # Update training metric.\n",
    "        train_acc_metric.update_state(y_batch_train, logits)\n",
    "\n",
    "        # Log every 5000 batches.\n",
    "        if step % 5000 == 0:\n",
    "            print(\n",
    "                \"Training loss (for one batch) at step %d: %.4f\"\n",
    "                % (step, float(loss_value))\n",
    "            )\n",
    "            print(\"Seen so far: %d samples\" % ((step + 1) * batch_size))\n",
    "            model.save_weights(\"tmp/ckpt-tf/ckpt-1\")\n",
    "            \n",
    "            # Run a validation loop at the end of each epoch.\n",
    "            best_checkpoint = \"tmp/ckpt-tf/ckpt-1\"\n",
    "            vf_model = ArcFaceModel(size=input_size,\n",
    "                                 backbone_type=backbone_type,               \n",
    "                                 training=False)\n",
    "            vf_model.load_weights(best_checkpoint)\n",
    "\n",
    "            acc_lfw, best_th = perform_val(\n",
    "                embd_shape, batch_size, vf_model, lfw, lfw_issame,\n",
    "                is_ccrop=False)\n",
    "            val_acc = acc_lfw\n",
    "            print(\"Validation acc: %.4f\" % (float(val_acc),))\n",
    "            print(\"Time taken: %.2fs\" % (time.time() - start_time))\n",
    "            if best_val_acc < val_acc:\n",
    "                best_val_acc = val_acc\n",
    "                model.save_weights(\"tmp/ckpt-tf/ckpt-best\")               \n",
    "            \n",
    "            # Display metrics at the end of each epoch.\n",
    "            train_acc = train_acc_metric.result()\n",
    "            print(\"Training acc over epoch: %.4f\" % (float(train_acc),))\n",
    "            # Reset training metrics at the end of each epoch\n",
    "            train_acc_metric.reset_states()\n",
    "    model.save_weights(f'tmp/ckpt-tf/epoch-{epoch}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}