{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## About\n\nIn this notebook, I'll train a model based on EfficientNetB0, GeM pooling, and ArcFace. \n\nThe whole training pipeline is built with TensorFlow, and the training will be done on TPU.\n\nThis notebook is based on [stratified-tfrecords-training-pipeline](https://www.kaggle.com/ks2019/stratified-tfrecords-training-pipeline).","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19"}},{"cell_type":"code","source":"!pip install efficientnet tensorflow_addons > /dev/null","metadata":{"execution":{"iopub.status.busy":"2021-08-31T00:23:54.158889Z","iopub.execute_input":"2021-08-31T00:23:54.159411Z","iopub.status.idle":"2021-08-31T00:24:03.496339Z","shell.execute_reply.started":"2021-08-31T00:23:54.159315Z","shell.execute_reply":"2021-08-31T00:24:03.495249Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport math\nimport random\nimport re\nimport warnings\nfrom pathlib import Path\nfrom typing import Optional, Tuple\n\nimport efficientnet.tfkeras as efn\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nimport tensorflow as tf\nimport tensorflow_addons as tfa\nfrom kaggle_datasets import KaggleDatasets\nfrom sklearn.model_selection import KFold","metadata":{"execution":{"iopub.status.busy":"2021-08-31T00:24:03.500320Z","iopub.execute_input":"2021-08-31T00:24:03.500620Z","iopub.status.idle":"2021-08-31T00:24:10.594119Z","shell.execute_reply.started":"2021-08-31T00:24:03.500589Z","shell.execute_reply":"2021-08-31T00:24:10.593183Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"tf.__version__","metadata":{"execution":{"iopub.status.busy":"2021-08-31T00:24:10.596468Z","iopub.execute_input":"2021-08-31T00:24:10.596882Z","iopub.status.idle":"2021-08-31T00:24:10.607582Z","shell.execute_reply.started":"2021-08-31T00:24:10.596839Z","shell.execute_reply":"2021-08-31T00:24:10.606744Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Config","metadata":{}},{"cell_type":"code","source":"NUM_FOLDS = 4\nIMAGE_SIZE = 256\nBATCH_SIZE = 64\nEFFICIENTNET_SIZE = 0\nWEIGHTS = \"imagenet\"\nN_CLASSES = 81313\nFOLDS = [0, 1, 2, 3]\nEPOCHS = 20\nSEED = 1213\n\nSAVEDIR = Path(\"./\")","metadata":{"execution":{"iopub.status.busy":"2021-08-31T00:28:24.619788Z","iopub.execute_input":"2021-08-31T00:28:24.620185Z","iopub.status.idle":"2021-08-31T00:28:24.626427Z","shell.execute_reply.started":"2021-08-31T00:28:24.620148Z","shell.execute_reply":"2021-08-31T00:28:24.625243Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Utilities","metadata":{}},{"cell_type":"code","source":"def set_seed(seed=42):\n    random.seed(seed)\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n    np.random.seed(seed)\n    tf.random.set_seed(seed)\n\n\nset_seed(SEED)","metadata":{"execution":{"iopub.status.busy":"2021-08-31T00:24:10.621795Z","iopub.execute_input":"2021-08-31T00:24:10.622264Z","iopub.status.idle":"2021-08-31T00:24:10.631339Z","shell.execute_reply.started":"2021-08-31T00:24:10.622235Z","shell.execute_reply":"2021-08-31T00:24:10.630655Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def auto_select_accelerator():\n    try:\n        tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n        tf.config.experimental_connect_to_cluster(tpu)\n        tf.tpu.experimental.initialize_tpu_system(tpu)\n        strategy = tf.distribute.experimental.TPUStrategy(tpu)\n        print(\"Running on TPU:\", tpu.master())\n    except ValueError:\n        strategy = tf.distribute.get_strategy()\n    print(f\"Running on {strategy.num_replicas_in_sync} replicas\")\n    return strategy","metadata":{"execution":{"iopub.status.busy":"2021-08-31T00:24:10.632730Z","iopub.execute_input":"2021-08-31T00:24:10.633145Z","iopub.status.idle":"2021-08-31T00:24:10.642409Z","shell.execute_reply.started":"2021-08-31T00:24:10.633106Z","shell.execute_reply":"2021-08-31T00:24:10.641715Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"strategy = auto_select_accelerator()\nREPLICAS = strategy.num_replicas_in_sync\nAUTO = tf.data.experimental.AUTOTUNE","metadata":{"execution":{"iopub.status.busy":"2021-08-31T00:24:10.643818Z","iopub.execute_input":"2021-08-31T00:24:10.644448Z","iopub.status.idle":"2021-08-31T00:24:16.308310Z","shell.execute_reply.started":"2021-08-31T00:24:10.644405Z","shell.execute_reply":"2021-08-31T00:24:16.307282Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Data Loading","metadata":{}},{"cell_type":"code","source":"gcs_paths = []\nfor i in range(5):\n    gcs_path = KaggleDatasets().get_gcs_path(f\"landmark-recognition-2021-tfrecords-fold{i}\")\n    print(gcs_path)\n    gcs_paths.append(gcs_path)\n    \nall_files = []\nfor path in gcs_paths:\n    all_files.extend(np.sort(np.array(tf.io.gfile.glob(path + \"/*.tfrec\"))))\n\nprint(\"train files: \", len(all_files))","metadata":{"execution":{"iopub.status.busy":"2021-08-31T00:25:25.487583Z","iopub.execute_input":"2021-08-31T00:25:25.488122Z","iopub.status.idle":"2021-08-31T00:25:27.554443Z","shell.execute_reply.started":"2021-08-31T00:25:25.488076Z","shell.execute_reply":"2021-08-31T00:25:27.553388Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def decode_image(image_data):\n    image = tf.image.decode_jpeg(image_data, channels=3)\n    image = tf.image.resize(image, size=(IMAGE_SIZE, IMAGE_SIZE))\n    image = tf.cast(image, tf.float32) / 255.0\n    return image\n\n\ndef read_labeled_tfrecord(example):\n    tfrec_format = {\n        \"image_name\": tf.io.FixedLenFeature([], tf.string),\n        \"image\": tf.io.FixedLenFeature([], tf.string),\n        \"target\": tf.io.FixedLenFeature([], tf.int64),\n    }\n\n    example = tf.io.parse_single_example(example, tfrec_format)\n    posting_id = example[\"image_name\"]\n    image = decode_image(example[\"image\"])\n    label_group = tf.cast(example[\"target\"], tf.int32)\n    matches = 1\n    return posting_id, image, label_group, matches\n\n\ndef arcface_format(posting_id, image, label_group, matches):\n    return posting_id, {'inp1': image, 'inp2': label_group}, label_group, matches\n\n\n# This function loads TF Records and parse them into tensors\ndef load_dataset(filenames, batch_size=64, cache=False, repeat=False, shuffle=False):        \n    dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads = AUTO)\n    if cache:\n        dataset = dataset.cache()\n\n    if shuffle:\n        dataset = dataset.shuffle(2048)\n        opt = tf.data.Options()\n        opt.experimental_deterministic = False\n        dataset = dataset.with_options(opt)\n\n    dataset = dataset.map(read_labeled_tfrecord, num_parallel_calls = AUTO) \n    dataset = dataset.map(arcface_format, num_parallel_calls=AUTO)\n    if repeat:\n        dataset = dataset.repeat()\n    dataset = dataset.batch(batch_size * REPLICAS)\n    dataset = dataset.prefetch(AUTO)\n    return dataset\n\n\ndef count_data_items(filenames):\n    # The number of data items is written in the name of the .tfrec files, i.e. flowers00-230.tfrec = 230 data items\n    n = [int(re.compile(r\"-([0-9]*)\\.\").search(filename).group(1)) for filename in filenames]\n    return np.sum(n)","metadata":{"execution":{"iopub.status.busy":"2021-08-31T00:25:51.686483Z","iopub.execute_input":"2021-08-31T00:25:51.686856Z","iopub.status.idle":"2021-08-31T00:25:51.701715Z","shell.execute_reply.started":"2021-08-31T00:25:51.686820Z","shell.execute_reply":"2021-08-31T00:25:51.700321Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"NUM_TRAINING_IMAGES = count_data_items(all_files)\nNUM_TRAINING_IMAGES","metadata":{"execution":{"iopub.status.busy":"2021-08-31T00:25:59.781935Z","iopub.execute_input":"2021-08-31T00:25:59.782339Z","iopub.status.idle":"2021-08-31T00:25:59.790312Z","shell.execute_reply.started":"2021-08-31T00:25:59.782302Z","shell.execute_reply":"2021-08-31T00:25:59.789258Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model","metadata":{}},{"cell_type":"code","source":"class GeM(tf.keras.layers.Layer):\n    def __init__(self, pool_size, init_norm=3.0, normalize=False, **kwargs):\n        self.pool_size = pool_size\n        self.init_norm = init_norm\n        self.normalize = normalize\n\n        super(GeM, self).__init__(**kwargs)\n\n    def get_config(self):\n        config = super().get_config().copy()\n        config.update({\n            'pool_size': self.pool_size,\n            'init_norm': self.init_norm,\n            'normalize': self.normalize,\n        })\n        return config\n\n    def build(self, input_shape):\n        feature_size = input_shape[-1]\n        self.p = self.add_weight(name='norms', shape=(feature_size,),\n                                 initializer=tf.keras.initializers.constant(self.init_norm),\n                                 trainable=True)\n        super(GeM, self).build(input_shape)\n\n    def call(self, inputs):\n        x = inputs\n        x = tf.math.maximum(x, 1e-6)\n        x = tf.pow(x, self.p)\n\n        x = tf.nn.avg_pool(x, self.pool_size, self.pool_size, 'VALID')\n        x = tf.pow(x, 1.0 / self.p)\n\n        if self.normalize:\n            x = tf.nn.l2_normalize(x, 1)\n        return x\n\n    def compute_output_shape(self, input_shape):\n        return tuple([None, input_shape[-1]])","metadata":{"execution":{"iopub.status.busy":"2021-08-31T00:26:23.774023Z","iopub.execute_input":"2021-08-31T00:26:23.774411Z","iopub.status.idle":"2021-08-31T00:26:23.786182Z","shell.execute_reply.started":"2021-08-31T00:26:23.774370Z","shell.execute_reply":"2021-08-31T00:26:23.784861Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ArcMarginProduct(tf.keras.layers.Layer):\n    '''\n    Implements large margin arc distance.\n\n    Reference:\n        https://arxiv.org/pdf/1801.07698.pdf\n        https://github.com/lyakaap/Landmark2019-1st-and-3rd-Place-Solution/\n            blob/master/src/modeling/metric_learning.py\n    '''\n    def __init__(self, n_classes, s=30, m=0.50, easy_margin=False,\n                 ls_eps=0.0, **kwargs):\n\n        super(ArcMarginProduct, self).__init__(**kwargs)\n\n        self.n_classes = n_classes\n        self.s = s\n        self.m = m\n        self.ls_eps = ls_eps\n        self.easy_margin = easy_margin\n        self.cos_m = tf.math.cos(m)\n        self.sin_m = tf.math.sin(m)\n        self.th = tf.math.cos(math.pi - m)\n        self.mm = tf.math.sin(math.pi - m) * m\n\n    def get_config(self):\n\n        config = super().get_config().copy()\n        config.update({\n            'n_classes': self.n_classes,\n            's': self.s,\n            'm': self.m,\n            'ls_eps': self.ls_eps,\n            'easy_margin': self.easy_margin,\n        })\n        return config\n\n    def build(self, input_shape):\n        super(ArcMarginProduct, self).build(input_shape[0])\n\n        self.W = self.add_weight(\n            name='W',\n            shape=(int(input_shape[0][-1]), self.n_classes),\n            initializer='glorot_uniform',\n            dtype='float32',\n            trainable=True,\n            regularizer=None)\n\n    def call(self, inputs):\n        X, y = inputs\n        y = tf.cast(y, dtype=tf.int32)\n        cosine = tf.matmul(\n            tf.math.l2_normalize(X, axis=1),\n            tf.math.l2_normalize(self.W, axis=0)\n        )\n        sine = tf.math.sqrt(1.0 - tf.math.pow(cosine, 2))\n        phi = cosine * self.cos_m - sine * self.sin_m\n        if self.easy_margin:\n            phi = tf.where(cosine > 0, phi, cosine)\n        else:\n            phi = tf.where(cosine > self.th, phi, cosine - self.mm)\n        one_hot = tf.cast(\n            tf.one_hot(y, depth=self.n_classes),\n            dtype=cosine.dtype\n        )\n        if self.ls_eps > 0:\n            one_hot = (1 - self.ls_eps) * one_hot + self.ls_eps / self.n_classes\n\n        output = (one_hot * phi) + ((1.0 - one_hot) * cosine)\n        output *= self.s\n        return output","metadata":{"execution":{"iopub.status.busy":"2021-08-31T00:26:35.893648Z","iopub.execute_input":"2021-08-31T00:26:35.894063Z","iopub.status.idle":"2021-08-31T00:26:35.912153Z","shell.execute_reply.started":"2021-08-31T00:26:35.894024Z","shell.execute_reply":"2021-08-31T00:26:35.910680Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def build_model(size=256, efficientnet_size=0, weights=\"imagenet\", count=0):\n    inp = tf.keras.layers.Input(shape=(size, size, 3), name=\"inp1\")\n    label = tf.keras.layers.Input(shape=(), name=\"inp2\")\n    x = getattr(efn, f\"EfficientNetB{efficientnet_size}\")(\n        weights=weights, include_top=False, input_shape=(size, size, 3))(inp)\n    x = GeM(8)(x)\n    x = tf.keras.layers.Flatten()(x)\n    x = tf.keras.layers.Dense(512, name=\"dense_before_arcface\", kernel_initializer=\"he_normal\")(x)\n    x = tf.keras.layers.BatchNormalization()(x)\n    x = ArcMarginProduct(\n        n_classes=N_CLASSES,\n        s=30,\n        m=0.5,\n        name=\"head/arc_margin\",\n        dtype=\"float32\"\n    )([x, label])\n    output = tf.keras.layers.Softmax(dtype=\"float32\")(x)\n    model = tf.keras.Model(inputs=[inp, label], outputs=[output])\n    lr_decayed_fn = tf.keras.experimental.CosineDecay(1e-3, count)\n    opt = tfa.optimizers.AdamW(lr_decayed_fn, learning_rate=1e-4)\n    model.compile(\n        optimizer=opt,\n        loss=[tf.keras.losses.SparseCategoricalCrossentropy()],\n        metrics=[tf.keras.metrics.SparseCategoricalAccuracy()]\n    )\n    return model","metadata":{"execution":{"iopub.status.busy":"2021-08-31T00:37:39.938076Z","iopub.execute_input":"2021-08-31T00:37:39.938436Z","iopub.status.idle":"2021-08-31T00:37:39.951089Z","shell.execute_reply.started":"2021-08-31T00:37:39.938407Z","shell.execute_reply":"2021-08-31T00:37:39.950110Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Other training utilities","metadata":{}},{"cell_type":"code","source":"def get_lr_callback(plot=False):\n    lr_start   = 1e-3\n    lr_max     = 0.00003 * BATCH_SIZE  \n    lr_min     = 1e-5\n    lr_ramp_ep = 4\n    lr_sus_ep  = 0\n    lr_decay   = 0.9\n   \n    def lrfn(epoch):\n        if epoch < lr_ramp_ep:\n            lr = (lr_max - lr_start) / lr_ramp_ep * epoch + lr_start\n            \n        elif epoch < lr_ramp_ep + lr_sus_ep:\n            lr = lr_max\n            \n        else:\n            lr = (lr_max - lr_min) * lr_decay**(epoch - lr_ramp_ep - lr_sus_ep) + lr_min\n            \n        return lr\n        \n    if plot:\n        epochs = list(range(EPOCHS))\n        learning_rates = [lrfn(x) for x in epochs]\n        plt.scatter(epochs,learning_rates)\n        plt.show()\n\n    lr_callback = tf.keras.callbacks.LearningRateScheduler(lrfn, verbose=False)\n    return lr_callback\n\nget_lr_callback(plot=True)","metadata":{"execution":{"iopub.status.busy":"2021-08-31T00:37:41.587051Z","iopub.execute_input":"2021-08-31T00:37:41.587634Z","iopub.status.idle":"2021-08-31T00:37:41.755727Z","shell.execute_reply.started":"2021-08-31T00:37:41.587600Z","shell.execute_reply":"2021-08-31T00:37:41.754954Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training","metadata":{}},{"cell_type":"code","source":"kf = KFold(n_splits=NUM_FOLDS, shuffle=True, random_state=SEED)\nfiles_train_all = np.array(all_files)","metadata":{"execution":{"iopub.status.busy":"2021-08-31T00:37:43.313409Z","iopub.execute_input":"2021-08-31T00:37:43.314017Z","iopub.status.idle":"2021-08-31T00:37:43.319050Z","shell.execute_reply.started":"2021-08-31T00:37:43.313971Z","shell.execute_reply":"2021-08-31T00:37:43.317983Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for fold, (trn_idx, val_idx) in enumerate(kf.split(files_train_all)):\n    if fold not in FOLDS:\n        continue\n    files_train = files_train_all[trn_idx]\n    files_valid = files_train_all[val_idx]\n\n    print(\"=\" * 120)\n    print(f\"Fold {fold}\")\n    print(\"=\" * 120)\n\n    train_image_count = count_data_items(files_train)\n    valid_image_count = count_data_items(files_valid)\n\n    tf.keras.backend.clear_session()\n    strategy = auto_select_accelerator()\n\n    with strategy.scope():\n        model = build_model(\n            size=IMAGE_SIZE,\n            efficientnet_size=EFFICIENTNET_SIZE,\n            weights=WEIGHTS,\n            count=train_image_count // BATCH_SIZE // REPLICAS // 4\n        )\n\n    model_ckpt = tf.keras.callbacks.ModelCheckpoint(\n        str(SAVEDIR / f\"fold{fold}.h5\"), monitor=\"val_loss\", verbose=1, save_best_only=True,\n        save_weights_only=True, mode=\"min\", save_freq=\"epoch\"\n    )\n\n    train_dataset = load_dataset(files_train, batch_size=BATCH_SIZE, shuffle=True, repeat=True)\n    train_dataset = train_dataset.map(lambda posting_id, image, label_group, matches: (image, label_group))\n\n    valid_dataset = load_dataset(files_valid, batch_size=BATCH_SIZE * 2, shuffle=False, repeat=False)\n    valid_dataset = valid_dataset.map(lambda posting_id, image, label_group, matches: (image, label_group))\n\n    STEPS_PER_EPOCH = train_image_count // BATCH_SIZE // REPLICAS // 4\n    history = model.fit(\n        train_dataset,\n        epochs=EPOCHS,\n        callbacks=[model_ckpt, get_lr_callback()],\n        steps_per_epoch=STEPS_PER_EPOCH,\n        validation_data=valid_dataset,\n        verbose=1\n    )","metadata":{"execution":{"iopub.status.busy":"2021-08-31T00:37:49.494246Z","iopub.execute_input":"2021-08-31T00:37:49.494660Z","iopub.status.idle":"2021-08-31T00:43:44.716918Z","shell.execute_reply.started":"2021-08-31T00:37:49.494596Z","shell.execute_reply":"2021-08-31T00:43:44.714402Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}