{"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":"Hello fellow Kagglers,\n\nThis notebook demonstrates how to train a CNN using ArcFace loss on the Google Landmark 2021 dataset. The EfficientNetV2-S CNN is used, introduced in [this](https://arxiv.org/pdf/2104.00298.pdf) paper. Moreover, the ArcFace loss is used to create class embeddings which are close to eachother, the ArcFace paper can be found [here](https://arxiv.org/pdf/1801.07698.pdf).\n\nSince the given dataset is huge with ~1.5M images training efficiency is key. To optimize training the images are converted to TFRrecords for fast reading, these datasets can be found here: [Part 1](https://www.kaggle.com/markwijkhuizen/landmark-recognition-2021-tfrecords-384-part-1), [Part 2](https://www.kaggle.com/markwijkhuizen/landmark-recognition-2021-tfrecords-384-part-2), [Part 3](https://www.kaggle.com/markwijkhuizen/landmark-recognition-2021-tfrecords-384-part-3). The last trick was to use bfloat16 training, which is a 16 bits float with a lower precision than a conventional 16 bits float, but the range of a 32 bits float. This reduces the computation time and allows for bigger batch sizes.","metadata":{}},{"cell_type":"code","source":"import warnings\nwarnings.simplefilter('ignore')","metadata":{"execution":{"iopub.status.busy":"2021-08-31T09:16:59.271944Z","iopub.execute_input":"2021-08-31T09:16:59.272383Z","iopub.status.idle":"2021-08-31T09:16:59.284523Z","shell.execute_reply.started":"2021-08-31T09:16:59.272295Z","shell.execute_reply":"2021-08-31T09:16:59.283415Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"This next line adds the [EfficientNetV2 GitHub repository](https://github.com/google/automl/tree/master/efficientnetv2) with the corresponding weights to the system.","metadata":{}},{"cell_type":"code","source":"import sys\nsys.path.append('/kaggle/input/efficientnetv2-pretrained-imagenet21k-weights')","metadata":{"execution":{"iopub.status.busy":"2021-08-31T09:16:59.307785Z","iopub.execute_input":"2021-08-31T09:16:59.308710Z","iopub.status.idle":"2021-08-31T09:16:59.312868Z","shell.execute_reply.started":"2021-08-31T09:16:59.308661Z","shell.execute_reply":"2021-08-31T09:16:59.311962Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport tensorflow as tf\nimport matplotlib.pyplot as plt\n\nfrom tensorflow.keras.mixed_precision import experimental as mixed_precision\nfrom kaggle_datasets import KaggleDatasets\nfrom tqdm.notebook import tqdm\nfrom brain_automl.efficientnetv2 import effnetv2_model, preprocessing\n\nimport re\nimport os\nimport io\nimport time\nimport pickle\nimport math\nimport random\nimport sys\n\nprint(f'tensorflow version: {tf.__version__}')\nprint(f'tensorflow keras version: {tf.keras.__version__}')\nprint(f'python version: P{sys.version}')","metadata":{"execution":{"iopub.status.busy":"2021-08-31T09:16:59.353554Z","iopub.execute_input":"2021-08-31T09:16:59.354424Z","iopub.status.idle":"2021-08-31T09:17:06.006863Z","shell.execute_reply.started":"2021-08-31T09:16:59.354360Z","shell.execute_reply":"2021-08-31T09:17:06.005729Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Seed all random number generators\ndef seed_everything(seed):\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    random.seed(seed)\n    np.random.seed(seed)\n    tf.random.set_seed(seed)\n    \nseed_everything(42)","metadata":{"execution":{"iopub.status.busy":"2021-08-31T09:17:06.008444Z","iopub.execute_input":"2021-08-31T09:17:06.008763Z","iopub.status.idle":"2021-08-31T09:17:06.014968Z","shell.execute_reply.started":"2021-08-31T09:17:06.008734Z","shell.execute_reply":"2021-08-31T09:17:06.013576Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Detect hardware, return appropriate distribution strategy\ntry:\n    TPU = tf.distribute.cluster_resolver.TPUClusterResolver()  # TPU detection. No parameters necessary if TPU_NAME environment variable is set. On Kaggle this is always the case.\n    print('Running on TPU ', TPU.master())\nexcept ValueError:\n    print('Running on GPU')\n    TPU = None\n\nif TPU:\n    tf.config.experimental_connect_to_cluster(TPU)\n    tf.tpu.experimental.initialize_tpu_system(TPU)\n    strategy = tf.distribute.experimental.TPUStrategy(TPU)\nelse:\n    strategy = tf.distribute.get_strategy() # default distribution strategy in Tensorflow. Works on CPU and single GPU.\n\nREPLICAS = strategy.num_replicas_in_sync\nprint(f'REPLICAS: {REPLICAS}')","metadata":{"execution":{"iopub.status.busy":"2021-08-31T09:17:06.017420Z","iopub.execute_input":"2021-08-31T09:17:06.017900Z","iopub.status.idle":"2021-08-31T09:17:11.924743Z","shell.execute_reply.started":"2021-08-31T09:17:06.017862Z","shell.execute_reply":"2021-08-31T09:17:11.923673Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DEBUG = False\n\n# Image dimensions\nIMG_SIZE = 384\nN_CHANNELS = 3\nINPUT_SHAPE = (IMG_SIZE, IMG_SIZE, N_CHANNELS)\n# Dataset size\nN_SAMPLES = 1580470\n\n# EfficientNet version, s, l, xl, xxl\nEFN_SIZE = 's'\n\n# Batch size per replica, there are 8 replicas resulting in a batch size of 1024\nBATCH_SIZE_BASE = 6 if DEBUG else (128 if TPU else 16)\nBATCH_SIZE = BATCH_SIZE_BASE * REPLICAS\n\nMODEL_POLICY = 'mixed_bfloat16' # float32 or mixed_bfloat16\nIMAGE_DTYPE = tf.bfloat16 if MODEL_POLICY == 'mixed_bfloat16' else tf.float32\nLABEL_DTYPE = tf.int32\n\n# Imagenet mean and standard deviation for normalizing images\nIMAGENET_MEAN = tf.constant([0.485, 0.456, 0.406], dtype=tf.float32)\nIMAGENET_STD = tf.constant([0.229, 0.224, 0.225], dtype=tf.float32)\n\n# Tensorflow AUTO flag\nAUTO = tf.data.experimental.AUTOTUNE\n\nprint(f'BATCH_SIZE: {BATCH_SIZE}, IMAGE_DTYPE: {IMAGE_DTYPE}, LABEL_DTYPE: {LABEL_DTYPE}')\nprint(f'MODEL_POLICY: {MODEL_POLICY}')","metadata":{"execution":{"iopub.status.busy":"2021-08-31T09:17:11.926274Z","iopub.execute_input":"2021-08-31T09:17:11.926726Z","iopub.status.idle":"2021-08-31T09:17:11.936473Z","shell.execute_reply.started":"2021-08-31T09:17:11.926680Z","shell.execute_reply":"2021-08-31T09:17:11.935417Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Landmark\\_id to label mappers","metadata":{}},{"cell_type":"code","source":"# mappers from landmark_id to label and vice versa\nwith open('/kaggle/input/landmark-recognition-2021-tfrecords-384-part-1/label2landmark_id.pkl', 'rb') as f:\n    label2landmark_id = pickle.load(f)\n    \nwith open('/kaggle/input/landmark-recognition-2021-tfrecords-384-part-1/landmark_id2label.pkl', 'rb') as f:\n    landmark_id2label = pickle.load(f)","metadata":{"execution":{"iopub.status.busy":"2021-08-31T09:17:11.937838Z","iopub.execute_input":"2021-08-31T09:17:11.938295Z","iopub.status.idle":"2021-08-31T09:17:12.043325Z","shell.execute_reply.started":"2021-08-31T09:17:11.938250Z","shell.execute_reply":"2021-08-31T09:17:12.042219Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"N_CLASSES = len(label2landmark_id.keys())\nprint(f'N_CLASSES: {N_CLASSES}')","metadata":{"execution":{"iopub.status.busy":"2021-08-31T09:17:12.045016Z","iopub.execute_input":"2021-08-31T09:17:12.045441Z","iopub.status.idle":"2021-08-31T09:17:12.051328Z","shell.execute_reply.started":"2021-08-31T09:17:12.045395Z","shell.execute_reply":"2021-08-31T09:17:12.050226Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset","metadata":{}},{"cell_type":"markdown","source":"Decode function for the TFRecords. Documentation on TFRecords can be found [here](https://www.tensorflow.org/api_docs/python/tf/data/TFRecordDataset). Each TFRecords contains 3000 images which can be processed in one go. This is much faster than reading all images one by one.","metadata":{}},{"cell_type":"code","source":"def decode_tfrecord(record_bytes):\n    # Data the sample contains\n    features = tf.io.parse_single_example(record_bytes, {\n        'image': tf.io.FixedLenFeature([], tf.string),\n        'label': tf.io.FixedLenFeature([], tf.int64),\n        'width': tf.io.FixedLenFeature([], tf.int64),\n        'height': tf.io.FixedLenFeature([], tf.int64),\n    })\n\n    image = tf.io.decode_jpeg(features['image'])\n    label = tf.cast(features['label'], dtype=tf.int32)\n    height = features['height']\n    width = features['width']\n    \n    # Cutout Random Square\n    if height != width:\n        if height > width:\n            # Get random offset\n            offset = tf.random.uniform(shape=(), minval=0, maxval=height-width, dtype=tf.int64)\n            image = tf.slice(image, [offset, 0, 0], [width, width, N_CHANNELS])\n        else:\n            # Get random offset\n            offset = tf.random.uniform(shape=(), minval=0, maxval=width-height, dtype=tf.int64)\n            image = tf.slice(image, [0, offset, 0], [height, height, N_CHANNELS])\n    \n    # Reshape and Normalize\n    size = tf.math.reduce_min([height, width])\n    image = tf.reshape(image, [size, size, N_CHANNELS])\n    image = tf.image.resize(image, [IMG_SIZE, IMG_SIZE])\n    image = tf.cast(image, tf.float32)  / 255.0\n    image = (image - IMAGENET_MEAN) / IMAGENET_STD\n    \n    # One hot encode the label, required for metrics\n    label_one_hot = tf.one_hot(label, N_CLASSES, dtype=tf.uint8)\n    \n    # Cast image if not desired dtype, will be converted to bfloat16 on TPU\n    if image.dtype != IMAGE_DTYPE:\n        image = tf.cast(image, IMAGE_DTYPE)\n    \n    # CNN required both image and label, passed as dictionary\n    return { 'image': image, 'label': label }, label_one_hot","metadata":{"execution":{"iopub.status.busy":"2021-08-31T09:17:12.052791Z","iopub.execute_input":"2021-08-31T09:17:12.053222Z","iopub.status.idle":"2021-08-31T09:17:12.069454Z","shell.execute_reply.started":"2021-08-31T09:17:12.053177Z","shell.execute_reply":"2021-08-31T09:17:12.068203Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Simple function to benchmark the dataset, images will be read with ~6000 images/second!\ndef benchmark_dataset(dataset, num_epochs=3, n_steps_per_epoch=25, bs=BATCH_SIZE):\n    start_time = time.perf_counter()\n    for epoch_num in range(num_epochs):\n        for idx, (inputs, labels) in enumerate(dataset.take(n_steps_per_epoch + 1)):\n            images = inputs['image']\n            if idx == 0:\n                epoch_start = time.perf_counter()\n            elif idx == 1 and epoch_num == 0:\n                print(f'image shape: {images.shape}, image dtype: {images.dtype}')\n            else:\n                pass\n        epoch_t = time.perf_counter() - epoch_start\n        mean_step_t = round(epoch_t / n_steps_per_epoch * 1000, 1)\n        n_imgs_per_s = int(1 / (mean_step_t / 1000) * bs)\n        print(f'epoch {epoch_num} took: {round(epoch_t, 2)} sec, mean step duration: {mean_step_t}ms, images/s: {n_imgs_per_s}')","metadata":{"execution":{"iopub.status.busy":"2021-08-31T09:17:12.072658Z","iopub.execute_input":"2021-08-31T09:17:12.073005Z","iopub.status.idle":"2021-08-31T09:17:12.085760Z","shell.execute_reply.started":"2021-08-31T09:17:12.072957Z","shell.execute_reply":"2021-08-31T09:17:12.084812Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Plots a batch of images\ndef show_batch(dataset, rows=4, cols=3):\n    inputs, lbls = next(iter(dataset))\n    imgs = inputs['image']\n    fig, axes = plt.subplots(nrows=rows, ncols=cols, figsize=(cols*7, rows*4))\n    for r in range(rows):\n        for c in range(cols):\n            img = imgs[r*cols+c].numpy().astype(np.float32)\n            img += abs(img.min())\n            img /= img.max()\n            axes[r, c].imshow(img)\n            axes[r, c].set_title(f'Label: {np.argmax(lbls[r*cols+c])}')","metadata":{"execution":{"iopub.status.busy":"2021-08-31T09:17:12.088340Z","iopub.execute_input":"2021-08-31T09:17:12.088771Z","iopub.status.idle":"2021-08-31T09:17:12.102862Z","shell.execute_reply.started":"2021-08-31T09:17:12.088728Z","shell.execute_reply":"2021-08-31T09:17:12.101815Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train Dataset","metadata":{}},{"cell_type":"code","source":"# Google Cloud paths to the TFRecords datasets, required as TPU will read from Google Cloud only\nGCS_DS_PATH_1 = KaggleDatasets().get_gcs_path('landmark-recognition-2021-tfrecords-384-part-1')\nGCS_DS_PATH_2 = KaggleDatasets().get_gcs_path('landmark-recognition-2021-tfrecords-384-part-2')\nGCS_DS_PATH_3 = KaggleDatasets().get_gcs_path('landmark-recognition-2021-tfrecords-384-part-3')","metadata":{"execution":{"iopub.status.busy":"2021-08-31T09:17:12.105016Z","iopub.execute_input":"2021-08-31T09:17:12.105752Z","iopub.status.idle":"2021-08-31T09:17:13.201237Z","shell.execute_reply.started":"2021-08-31T09:17:12.105713Z","shell.execute_reply":"2021-08-31T09:17:13.200275Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_train_dataset(bs=BATCH_SIZE):\n    # Ignore order, improves performance\n    ignore_order = tf.data.Options()\n    ignore_order.experimental_deterministic = False\n    \n    # Use glob to find all TFRecords files\n    FNAMES_TRAIN_TFRECORDS = (\n        tf.io.gfile.glob(f'{GCS_DS_PATH_1}/*.tfrecords') +\n        tf.io.gfile.glob(f'{GCS_DS_PATH_2}/*.tfrecords') +\n        tf.io.gfile.glob(f'{GCS_DS_PATH_3}/*.tfrecords')\n    )\n    \n    print(f'Found roughly {len(FNAMES_TRAIN_TFRECORDS) * int(3e3)} images, N_SAMPLES: {N_SAMPLES}')\n    \n    train_dataset = tf.data.TFRecordDataset(FNAMES_TRAIN_TFRECORDS, num_parallel_reads=AUTO)\n    train_dataset = train_dataset.with_options(ignore_order)\n    train_dataset = train_dataset.repeat()\n    train_dataset = train_dataset.map(decode_tfrecord, num_parallel_calls=AUTO)\n    train_dataset = train_dataset.batch(BATCH_SIZE)\n    train_dataset = train_dataset.prefetch(AUTO)\n    \n    return train_dataset\n\ntrain_dataset = get_train_dataset()","metadata":{"execution":{"iopub.status.busy":"2021-08-31T09:17:13.202626Z","iopub.execute_input":"2021-08-31T09:17:13.203045Z","iopub.status.idle":"2021-08-31T09:17:13.899621Z","shell.execute_reply.started":"2021-08-31T09:17:13.203003Z","shell.execute_reply":"2021-08-31T09:17:13.898542Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Benchmark the dataset, close to 6000 images/second can be read using TFRecords!\nbenchmark_dataset(train_dataset)","metadata":{"execution":{"iopub.status.busy":"2021-08-31T09:17:13.900881Z","iopub.execute_input":"2021-08-31T09:17:13.901206Z","iopub.status.idle":"2021-08-31T09:17:32.477918Z","shell.execute_reply.started":"2021-08-31T09:17:13.901177Z","shell.execute_reply":"2021-08-31T09:17:32.476843Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Sanity check, what type and shapes are the images and labels\ninputs, lbls_oh = next(iter(train_dataset))\nimgs = inputs['image']\nprint(f'imgs shape: {imgs.shape}, imgs dtype: {imgs.dtype}')\nimg0 = imgs[0].numpy().astype(np.float32)\ntrain_imgs_info = (img0.mean(), img0.std(), img0.min(), img0.max())\nprint('train img 0 mean: %.3f, 0 std: %.3f, min: %.3f, max: %.3f' % train_imgs_info)\n\n# Labels\nlbls = inputs['label'].numpy()\nprint(f'lbls shape: {lbls.shape}, lbls dtype: {lbls.dtype}')\nprint(f'lbls min: {lbls.min()}, lbls max: {lbls.max()}')\n\n# Labels One Hot Encoded\nlbls_oh_np = lbls_oh.numpy()\nprint(f'lbls_oh shape: {lbls_oh.shape}, lbls_oh dtype: {lbls_oh.dtype}')","metadata":{"execution":{"iopub.status.busy":"2021-08-31T09:17:32.479238Z","iopub.execute_input":"2021-08-31T09:17:32.479551Z","iopub.status.idle":"2021-08-31T09:17:33.878424Z","shell.execute_reply.started":"2021-08-31T09:17:32.479522Z","shell.execute_reply":"2021-08-31T09:17:33.877388Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"show_batch(train_dataset)","metadata":{"execution":{"iopub.status.busy":"2021-08-31T09:17:33.882041Z","iopub.execute_input":"2021-08-31T09:17:33.882376Z","iopub.status.idle":"2021-08-31T09:17:36.794517Z","shell.execute_reply.started":"2021-08-31T09:17:33.882346Z","shell.execute_reply":"2021-08-31T09:17:36.793296Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# ArcMargin Product","metadata":{}},{"cell_type":"markdown","source":"ArcMargin product used for the ArcFace loss. I can't explain the math behind it, but the basic idea is to cluster the embeddings of samples belonging to the same class close together. The difference between conventional softmax and ArcFace loss is nicely illustrated in [this](https://www.kaggle.com/chankhavu/keras-layers-arcface-cosface-adacos) notebook. If someone could explain the math behind it that would be very welcome :)","metadata":{}},{"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, 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 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=tf.float32,\n            trainable=True,\n            regularizer=None\n        )\n\n    def call(self, inputs):\n        X, y = inputs\n        \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        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        \n        return output","metadata":{"execution":{"iopub.status.busy":"2021-08-31T09:17:36.795935Z","iopub.execute_input":"2021-08-31T09:17:36.796496Z","iopub.status.idle":"2021-08-31T09:17:36.815734Z","shell.execute_reply.started":"2021-08-31T09:17:36.796456Z","shell.execute_reply":"2021-08-31T09:17:36.814635Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"code","source":"# EfficientNetV2-S pretrained weights also need to be read from Google Cloud\nGCS_WEIGHTS_PATH = KaggleDatasets().get_gcs_path('efficientnetv2-pretrained-imagenet21k-weights')","metadata":{"execution":{"iopub.status.busy":"2021-08-31T09:17:36.817032Z","iopub.execute_input":"2021-08-31T09:17:36.817389Z","iopub.status.idle":"2021-08-31T09:17:37.189115Z","shell.execute_reply.started":"2021-08-31T09:17:36.817352Z","shell.execute_reply":"2021-08-31T09:17:37.188197Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_model():\n    tf.keras.backend.clear_session()\n    # enable XLA optmizations\n    tf.config.optimizer.set_jit(True)\n    \n    # set half precision policy\n    mixed_precision.set_policy(MODEL_POLICY)\n\n    # Print compute and variable dtype, on TPU this will be bfloat16 for compute and float32 for variable\n    print(f'Compute dtype: {mixed_precision.global_policy().compute_dtype}')\n    print(f'Variable dtype: {mixed_precision.global_policy().variable_dtype}')\n\n    with strategy.scope():\n        WEIGHT_PATH = f'{GCS_WEIGHTS_PATH}/efficientnetv2-{EFN_SIZE}-21k-ft1k'\n        cnn = effnetv2_model.EffNetV2Model(model_name='efficientnetv2-s')\n        \n        # Inputs, note the names are equal to the dictionary keys in the dataset\n        image = tf.keras.layers.Input(INPUT_SHAPE, name='image', dtype=IMAGE_DTYPE)\n        label = tf.keras.layers.Input([], name='label', dtype=tf.int32)\n\n        # Build the model with a dummy call, this is required\n        cnn(tf.ones([1,*INPUT_SHAPE]), training=False)\n\n        # Get the latest checkpoint from path\n        ckpt = tf.train.latest_checkpoint(WEIGHT_PATH)\n\n        # Load the weights\n        cnn.load_weights(ckpt)\n        \n        # CNN call, we need only the output layer\n        x = cnn(image, features_only=True)[0]\n        # Global Average Pooling, cast to float32 for ArcMargin product\n        x = tf.keras.layers.GlobalAveragePooling2D(name='pooling', dtype=tf.float32)(x)\n        # Optional Dropout layer\n        x = tf.keras.layers.Dropout(0.00, name='dropout', dtype=tf.float32)(x)\n        # ArcMargin product\n        output = ArcMarginProduct(n_classes=N_CLASSES, name='arc_margin', dtype=tf.float32)([x, label])\n\n        # We will use the famous Adam optimizer for fast learning\n        optimizer = tf.keras.optimizers.Adam()\n\n        # Categorical Cross Entropy loss, from_logits=True so no softmax needed\n        loss = tf.keras.losses.CategoricalCrossentropy(from_logits=True)\n\n        # To track learning progress, top 1/10/100/1000 accuracies will be kept track of\n        metrics = [\n            tf.keras.metrics.TopKCategoricalAccuracy(k=1, name='accuracy@1'),\n            tf.keras.metrics.TopKCategoricalAccuracy(k=10, name='accuracy@10'),\n            tf.keras.metrics.TopKCategoricalAccuracy(k=100, name='accuracy@100'),\n            tf.keras.metrics.TopKCategoricalAccuracy(k=1000, name='accuracy@1000'),\n        ]\n\n        model = tf.keras.models.Model(inputs = [image, label], outputs = [output])\n        model.compile(optimizer=optimizer, loss=loss, metrics=metrics)\n\n        return model","metadata":{"execution":{"iopub.status.busy":"2021-08-31T09:17:37.190466Z","iopub.execute_input":"2021-08-31T09:17:37.190775Z","iopub.status.idle":"2021-08-31T09:17:37.204945Z","shell.execute_reply.started":"2021-08-31T09:17:37.190745Z","shell.execute_reply":"2021-08-31T09:17:37.203691Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = get_model()","metadata":{"execution":{"iopub.status.busy":"2021-08-31T09:17:37.206540Z","iopub.execute_input":"2021-08-31T09:17:37.207184Z","iopub.status.idle":"2021-08-31T09:18:51.735727Z","shell.execute_reply.started":"2021-08-31T09:17:37.206948Z","shell.execute_reply":"2021-08-31T09:18:51.734623Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Plot model summary\nmodel.summary()","metadata":{"execution":{"iopub.status.busy":"2021-08-31T09:18:51.736998Z","iopub.execute_input":"2021-08-31T09:18:51.737296Z","iopub.status.idle":"2021-08-31T09:18:51.797942Z","shell.execute_reply.started":"2021-08-31T09:18:51.737268Z","shell.execute_reply":"2021-08-31T09:18:51.796800Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Plot slightly more detailed model summary\ntf.keras.utils.plot_model(model, show_shapes=True, show_dtype=True, show_layer_names=True, expand_nested=False)","metadata":{"execution":{"iopub.status.busy":"2021-08-31T09:18:51.799225Z","iopub.execute_input":"2021-08-31T09:18:51.799524Z","iopub.status.idle":"2021-08-31T09:18:52.308523Z","shell.execute_reply.started":"2021-08-31T09:18:51.799496Z","shell.execute_reply":"2021-08-31T09:18:52.307314Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Learning Rate Scheduler","metadata":{}},{"cell_type":"markdown","source":"Because of the transfer learning approach, an exponential warmup is used with a cosine decay","metadata":{}},{"cell_type":"code","source":"# Due to the huge batch size of 1024 and usage of bfloat16 15 epochs are possible in a single run!\nEPOCHS = 2 if DEBUG else 15","metadata":{"execution":{"iopub.status.busy":"2021-08-31T09:18:52.311399Z","iopub.execute_input":"2021-08-31T09:18:52.311753Z","iopub.status.idle":"2021-08-31T09:18:52.316193Z","shell.execute_reply.started":"2021-08-31T09:18:52.311718Z","shell.execute_reply":"2021-08-31T09:18:52.315169Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# returns the learning rate given an epoch number\ndef lrfn(epoch):\n    # Config\n    LR_START = 1e-5 # start of learning rate\n    LR_MAX = 2e-4 # peak learning rate\n    LR_FINAL = 2e-5 # final learning rate\n    LR_RAMPUP_EPOCHS = 3 # number of exponential warmup epochs\n    LR_SUSTAIN_EPOCHS = 2 # number of epochs at maximum learning rate\n    \n    DECAY_EPOCHS = EPOCHS  - LR_RAMPUP_EPOCHS - LR_SUSTAIN_EPOCHS - 1\n    LR_EXP_DECAY = (LR_FINAL / LR_MAX) ** (1 / (EPOCHS - LR_RAMPUP_EPOCHS - LR_SUSTAIN_EPOCHS - 1))\n\n    if epoch < LR_RAMPUP_EPOCHS: # exponential warmup\n        lr = LR_START + (LR_MAX + LR_START) * (epoch / LR_RAMPUP_EPOCHS) ** 2.5\n    elif epoch < LR_RAMPUP_EPOCHS + LR_SUSTAIN_EPOCHS: # sustain lr\n        lr = LR_MAX\n    else: # cosine decay\n        epoch_diff = epoch - LR_RAMPUP_EPOCHS - LR_SUSTAIN_EPOCHS + 1\n        decay_factor = (epoch_diff / (DECAY_EPOCHS + 1)) * math.pi\n        decay_factor= (tf.math.cos(decay_factor).numpy() + 1) / 2\n        lr = LR_FINAL + (LR_MAX - LR_FINAL) * decay_factor\n\n    return round(lr, 8)","metadata":{"execution":{"iopub.status.busy":"2021-08-31T09:18:52.317529Z","iopub.execute_input":"2021-08-31T09:18:52.318054Z","iopub.status.idle":"2021-08-31T09:18:52.331898Z","shell.execute_reply.started":"2021-08-31T09:18:52.318018Z","shell.execute_reply":"2021-08-31T09:18:52.330590Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Plots the learning rate schedule\ndef plot_lr_schedule(lr_schedule, name):\n    plt.figure(figsize=(15,8))\n    plt.plot(lr_schedule)\n    x = np.arange(EPOCHS)\n    x_axis_labels = list(map(str, np.arange(1, EPOCHS+1)))\n    plt.xticks(x, x_axis_labels) # set tick step to 1 and let x axis start at 1\n    schedule_info = f'start: {lr_schedule[0]}, max: {max(lr_schedule)}, final: {lr_schedule[-1]}'\n    plt.title(f'Step Learning Rate Schedule {name}, {schedule_info}', size=16)\n    plt.grid()\n    plt.show()\n\n# Learning rate for encoder\nLR_SCHEDULE = [lrfn(step) for step in range(EPOCHS)]\nplot_lr_schedule(LR_SCHEDULE, 'Ecnoder')","metadata":{"execution":{"iopub.status.busy":"2021-08-31T09:18:52.335655Z","iopub.execute_input":"2021-08-31T09:18:52.335976Z","iopub.status.idle":"2021-08-31T09:18:52.587727Z","shell.execute_reply.started":"2021-08-31T09:18:52.335947Z","shell.execute_reply":"2021-08-31T09:18:52.586854Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Callbacks","metadata":{}},{"cell_type":"code","source":"# Learning rate callback\nlr_callback = tf.keras.callbacks.LearningRateScheduler(lambda step: lrfn(step), verbose=1)\n# Model checkpoint, saves weights if train loss reduces\nmodel_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(\n    'model.h5', monitor='loss', verbose=1, save_best_only=True, save_weights_only=True\n)","metadata":{"execution":{"iopub.status.busy":"2021-08-31T09:18:52.589018Z","iopub.execute_input":"2021-08-31T09:18:52.589468Z","iopub.status.idle":"2021-08-31T09:18:52.593935Z","shell.execute_reply.started":"2021-08-31T09:18:52.589436Z","shell.execute_reply":"2021-08-31T09:18:52.593191Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training","metadata":{}},{"cell_type":"code","source":"# Because of the repeating dataset the amount of steps per epoch needs to be defined\nSTEPS_PER_EPOCH = N_SAMPLES // BATCH_SIZE\n\nprint(f'BATCH_SIZE: {BATCH_SIZE}, STEPS_PER_EPOCH: {STEPS_PER_EPOCH}')","metadata":{"execution":{"iopub.status.busy":"2021-08-31T09:18:52.595057Z","iopub.execute_input":"2021-08-31T09:18:52.595463Z","iopub.status.idle":"2021-08-31T09:18:52.610780Z","shell.execute_reply.started":"2021-08-31T09:18:52.595434Z","shell.execute_reply":"2021-08-31T09:18:52.609573Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Train the model, each epoch takes just ~23 minutes!\nhistory = model.fit(\n    train_dataset,\n    steps_per_epoch = STEPS_PER_EPOCH,\n    epochs = EPOCHS,\n    verbose = 2,\n    callbacks = [\n        lr_callback,\n        model_checkpoint_callback,\n    ],\n)","metadata":{"execution":{"iopub.status.busy":"2021-08-31T09:18:52.612454Z","iopub.execute_input":"2021-08-31T09:18:52.612992Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training History","metadata":{}},{"cell_type":"markdown","source":"Plot metric history during training","metadata":{}},{"cell_type":"code","source":"# Function to plot the metric history\ndef plot_history_metric(history, metric, f_best):\n    plt.figure(figsize=(15, 8))\n    N_EPOCHS = len(history.history['loss'])\n    x = [1, 5] + [10 + 5 * idx for idx in range((N_EPOCHS - 10) // 5 + 1)]\n    x_ticks = np.arange(1, N_EPOCHS+1)\n    # summarize history for accuracy\n    plt.plot(x_ticks, history.history[metric])\n    values = history.history[metric]\n    argmin = f_best(values)\n    plt.scatter(argmin + 1, values[argmin], color='red', s=50, marker='o')\n    \n    plt.title(f'Model {metric}', fontsize=24, pad=10)\n    plt.ylabel(metric, fontsize=20, labelpad=10)\n    plt.xlabel('epoch', fontsize=20, labelpad=10)\n    plt.tick_params(axis='x', labelsize=8)\n    plt.xticks(x, fontsize=16) # set tick step to 1 and let x axis start at 1\n    plt.yticks(fontsize=16)\n    plt.legend(['train'],  prop={'size': 18})\n    plt.grid()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_history_metric(history, 'loss', np.argmin)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_history_metric(history, 'accuracy@1', np.argmax)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_history_metric(history, 'accuracy@10', np.argmax)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_history_metric(history, 'accuracy@100', np.argmax)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_history_metric(history, 'accuracy@1000', np.argmax)","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}