{"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 the training process for the Happy Whale 2022 competition using the EfficientNetV2-XL model with a DOLG head and ArcFace classifier.\n\nThis model architecture is based on the 1st place solution of the [Google Landmark Recognition 2021](https://www.kaggle.com/c/landmark-recognition-2021) competition by [Christof Henkel](https://www.kaggle.com/christofhenkel) who poblished his solution written in Pytorch on [GitHuB](https://github.com/ChristofHenkel/kaggle-landmark-2021-1st-place/blob/main/models/ch_mdl_dolg_efficientnet.py).\n\nThis solution is a Tensorflow implementation of his model architecture with some tweaks. Among others, this solution uses conventional global average pooling instead of generalized mean pooling and does not have a batch normalization and PReLu layer in the head.\n\nThe implementation of Christof Henkel is based on DOLG: Single-Stage Image Retrieval with Deep Orthogonal Fusion of Local and Global Features ([paper](https://arxiv.org/pdf/2108.02927.pdf)).\n\nThis notebook should contribute to this competition by showing a Tensorflow implementation of DOLG and, to the best of my knowledge, setting a new benchmark of a single fold model.\n\n[Inference Notebook](https://www.kaggle.com/markwijkhuizen/happy-whale-2022-efficientnetv2-xl-dolg-inference)","metadata":{}},{"cell_type":"code","source":"import sys\nsys.path.append('/kaggle/input/efficientnetv2-pretrained-imagenet21k-weights/brain_automl/')\nsys.path.append('/kaggle/input/efficientnetv2-pretrained-imagenet21k-weights/brain_automl/efficientnetv2/')","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:26:37.195011Z","iopub.execute_input":"2022-03-20T10:26:37.195505Z","iopub.status.idle":"2022-03-20T10:26:37.218107Z","shell.execute_reply.started":"2022-03-20T10:26:37.195382Z","shell.execute_reply":"2022-03-20T10:26:37.217452Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import warnings\nwarnings.simplefilter('ignore')","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:26:37.452567Z","iopub.execute_input":"2022-03-20T10:26:37.453040Z","iopub.status.idle":"2022-03-20T10:26:37.456987Z","shell.execute_reply.started":"2022-03-20T10:26:37.452986Z","shell.execute_reply":"2022-03-20T10:26:37.456240Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport tensorflow as tf\nimport tensorflow_addons as tfa\nimport tensorflow.keras.backend as K\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\n\nimport re\nimport os\nimport io\nimport time\nimport pickle\nimport math\nimport random\nimport sys\nimport imageio\nimport effnetv2_model\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":"2022-03-20T10:26:37.647226Z","iopub.execute_input":"2022-03-20T10:26:37.647502Z","iopub.status.idle":"2022-03-20T10:26:44.178217Z","shell.execute_reply.started":"2022-03-20T10:26:37.647470Z","shell.execute_reply":"2022-03-20T10:26:44.177212Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def 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":"2022-03-20T10:26:44.180150Z","iopub.execute_input":"2022-03-20T10:26:44.180622Z","iopub.status.idle":"2022-03-20T10:26:44.187144Z","shell.execute_reply.started":"2022-03-20T10:26:44.180553Z","shell.execute_reply":"2022-03-20T10:26:44.186136Z"},"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":"2022-03-20T10:26:44.189713Z","iopub.execute_input":"2022-03-20T10:26:44.190502Z","iopub.status.idle":"2022-03-20T10:26:44.216954Z","shell.execute_reply.started":"2022-03-20T10:26:44.190456Z","shell.execute_reply":"2022-03-20T10:26:44.216051Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Load Train","metadata":{}},{"cell_type":"code","source":"train = pd.read_csv('/kaggle/input/happy-whale-and-dolphin/train.csv')\ndisplay(train.head())","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:41:25.570516Z","iopub.execute_input":"2022-03-20T10:41:25.571304Z","iopub.status.idle":"2022-03-20T10:41:25.654006Z","shell.execute_reply.started":"2022-03-20T10:41:25.571259Z","shell.execute_reply":"2022-03-20T10:41:25.653028Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DEBUG = False\n\n# Input Image Shape\nIMG_SIZE = 640\nN_CHANNELS = 3\nINPUT_SHAPE = (IMG_SIZE, IMG_SIZE, N_CHANNELS)\nN_SAMPLES = len(train)\n\nN_EPOCHS = 20\n\n# Model Configuration\nEFNV2_SIZE = 'xl'\nDOLG_SIZE = 1024\nEMBEDDING_SIZE = 2048\n\n# Due to huge amount of parameters batch size of 16 is not possible\nBATCH_SIZE_BASE = 12\nBATCH_SIZE = BATCH_SIZE_BASE * REPLICAS\n\nCROP = True\n\n# ImageNet Normalization\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\nAUTO = tf.data.experimental.AUTOTUNE\nEPS = tf.keras.backend.epsilon()\n\nprint(f'N_SAMPLES: {N_SAMPLES}, BATCH_SIZE: {BATCH_SIZE}')","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:41:26.436474Z","iopub.execute_input":"2022-03-20T10:41:26.436750Z","iopub.status.idle":"2022-03-20T10:41:26.445237Z","shell.execute_reply.started":"2022-03-20T10:41:26.436721Z","shell.execute_reply":"2022-03-20T10:41:26.444560Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Number of Labels","metadata":{}},{"cell_type":"code","source":"N_INDIVIDUAL_IDS = train['individual_id'].nunique()\nprint(f'N_INDIVIDUAL_IDS: {N_INDIVIDUAL_IDS}')","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:41:27.324700Z","iopub.execute_input":"2022-03-20T10:41:27.325541Z","iopub.status.idle":"2022-03-20T10:41:27.339631Z","shell.execute_reply.started":"2022-03-20T10:41:27.325503Z","shell.execute_reply":"2022-03-20T10:41:27.338697Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset\n\nThe amazing cropped [backfintfrecords](https://www.kaggle.com/datasets/jpbremer/backfintfrecords) by [Jan Bre](https://www.kaggle.com/jpbremer) is used.","metadata":{}},{"cell_type":"code","source":"# Function to Decode TFRecords and augment the image\ndef decode_tfrecord(record_bytes):\n    features = tf.io.parse_single_example(record_bytes, {\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    target = tf.cast(features['target'], tf.int32)\n    \n    image = tf.io.decode_jpeg(features['image'])\n    # Resize Image\n    image = tf.image.resize(image, [IMG_SIZE, IMG_SIZE])\n    # Explicit reshape needed for TPU, tell cimpiler dimensions of image\n    image = tf.reshape(image, INPUT_SHAPE)\n        \n    # Image Augmentations: retrieved from:\n    # https://www.kaggle.com/aikhmelnytskyy/happywhale-arcface-baseline-eff7-tpu-768-concat?scriptVersionId=88596800&cellId=15\n    image = tf.image.random_flip_left_right(image)\n    image = tf.image.random_hue(image, 0.01)\n    image = tf.image.random_saturation(image, 0.70, 1.30)\n    image = tf.image.random_contrast(image, 0.80, 1.20)\n    image = tf.image.random_brightness(image, 0.10)\n        \n    # ImageNet Normalization\n    image = tf.cast(image, tf.float32)  / 255.0\n    image = (image - IMAGENET_MEAN) / IMAGENET_STD\n    \n    return { 'image': image, 'individual_id_input': target }, { 'individual_id': target }","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:30:56.496828Z","iopub.execute_input":"2022-03-20T10:30:56.497134Z","iopub.status.idle":"2022-03-20T10:30:56.506821Z","shell.execute_reply.started":"2022-03-20T10:30:56.497099Z","shell.execute_reply":"2022-03-20T10:30:56.505815Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Simple Function to benchmark the dataset to make sure the data loader won't form a bottleneck\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":"2022-03-20T10:31:20.762108Z","iopub.execute_input":"2022-03-20T10:31:20.762703Z","iopub.status.idle":"2022-03-20T10:31:20.770222Z","shell.execute_reply.started":"2022-03-20T10:31:20.762654Z","shell.execute_reply":"2022-03-20T10:31:20.769471Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Function to show a batch of images\ndef show_batch(dataset, rows=5, cols=4):\n    inputs, lbls = next(iter(dataset))\n    imgs = inputs['image']\n    fig, axes = plt.subplots(nrows=rows, ncols=cols, figsize=(cols*6, rows*6))\n    for r in range(rows):\n        for c in range(cols):\n            idx = r*cols+c\n            img = imgs[idx].numpy().astype(np.float32)\n            img += abs(img.min())\n            img /= img.max()\n            axes[r, c].imshow(img)\n            individual_id = lbls['individual_id'][idx]\n            axes[r, c].set_title(f'individual_id: {individual_id}')","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:31:29.086244Z","iopub.execute_input":"2022-03-20T10:31:29.086527Z","iopub.status.idle":"2022-03-20T10:31:29.094098Z","shell.execute_reply.started":"2022-03-20T10:31:29.086496Z","shell.execute_reply":"2022-03-20T10:31:29.093420Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train Dataset","metadata":{}},{"cell_type":"code","source":"# For TPU's the dataset needs to be stored in Google Cloud\n# Retrieve the Google Cloud location of the dataset\nGCS_DS_PATH = KaggleDatasets().get_gcs_path('backfintfrecords')","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:32:01.650964Z","iopub.execute_input":"2022-03-20T10:32:01.651268Z","iopub.status.idle":"2022-03-20T10:32:01.984440Z","shell.execute_reply.started":"2022-03-20T10:32:01.651233Z","shell.execute_reply":"2022-03-20T10:32:01.983561Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Train Dataset\ndef get_train_dataset(bs=BATCH_SIZE, center_cutout=False, dr=True, sr=True):\n    FNAMES_TRAIN_TFRECORDS = tf.io.gfile.glob(f'{GCS_DS_PATH}/*train*.tfrec')\n    \n    # Shuffle TFRecords\n    random.shuffle(FNAMES_TRAIN_TFRECORDS)\n    \n    train_dataset = tf.data.TFRecordDataset(FNAMES_TRAIN_TFRECORDS, num_parallel_reads=AUTO)\n    \n    # Shuffle and Repeat dataset\n    if sr:\n        ignore_order = tf.data.Options()\n        ignore_order.experimental_deterministic = False\n        \n        train_dataset = train_dataset.with_options(ignore_order)\n        train_dataset = train_dataset.shuffle(N_SAMPLES if TPU else 1024)\n        train_dataset = train_dataset.repeat()\n        \n    map_fn = lambda e: decode_tfrecord(e)\n    train_dataset = train_dataset.map(map_fn, num_parallel_calls=AUTO)\n    train_dataset = train_dataset.batch(bs, drop_remainder=dr)\n    train_dataset = train_dataset.prefetch(AUTO)\n    \n    return train_dataset","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:33:03.425250Z","iopub.execute_input":"2022-03-20T10:33:03.425523Z","iopub.status.idle":"2022-03-20T10:33:03.434039Z","shell.execute_reply.started":"2022-03-20T10:33:03.425497Z","shell.execute_reply":"2022-03-20T10:33:03.433136Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"benchmark_dataset(get_train_dataset(center_cutout=True))","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:33:05.376659Z","iopub.execute_input":"2022-03-20T10:33:05.377090Z","iopub.status.idle":"2022-03-20T10:33:35.065671Z","shell.execute_reply.started":"2022-03-20T10:33:05.377048Z","shell.execute_reply":"2022-03-20T10:33:35.064549Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Input Image statistics, verify normalization\ninputs, labels = next(iter(get_train_dataset()))\nimgs = inputs['image']\nprint(f'inputs keys: {inputs.keys()}, labels keys: {labels.keys()}')\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)","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:33:35.067983Z","iopub.execute_input":"2022-03-20T10:33:35.068272Z","iopub.status.idle":"2022-03-20T10:33:38.309995Z","shell.execute_reply.started":"2022-03-20T10:33:35.068233Z","shell.execute_reply":"2022-03-20T10:33:38.309017Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"show_batch(get_train_dataset(bs=32))","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:33:38.311416Z","iopub.execute_input":"2022-03-20T10:33:38.312012Z","iopub.status.idle":"2022-03-20T10:33:48.923705Z","shell.execute_reply.started":"2022-03-20T10:33:38.311964Z","shell.execute_reply":"2022-03-20T10:33:48.923028Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Augmentation Test\n\nVisualization of augmentations","metadata":{}},{"cell_type":"code","source":"def get_demo_image():\n    demo_image = imageio.imread('/kaggle/input/happy-whale-and-dolphin/train_images/000562241d384d.jpg')\n    h, w = demo_image.shape[:2]\n    demo_image = demo_image[:, (w - h) // 2:(w - h) // 2 + h,:]\n    demo_image = tf.constant(demo_image)\n    demo_image_h = tf.constant(demo_image.shape[0], dtype=tf.int64)\n    demo_image_w = tf.constant(demo_image.shape[1], dtype=tf.int64)\n    \n    return demo_image, demo_image_h, demo_image_w\n\ndemo_image, demo_image_h, demo_image_w = get_demo_image()","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:37:46.135808Z","iopub.execute_input":"2022-03-20T10:37:46.136180Z","iopub.status.idle":"2022-03-20T10:37:46.359952Z","shell.execute_reply.started":"2022-03-20T10:37:46.136141Z","shell.execute_reply":"2022-03-20T10:37:46.359117Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def decode_tfrecord_demo(record_bytes):\n    image = demo_image\n    height = demo_image_h\n    width = demo_image_w\n    \n    # Resize to IMG_SIZE\n    image = tf.image.resize(image, [IMG_SIZE, IMG_SIZE], method=tf.image.ResizeMethod.BICUBIC)\n        \n    # Image Augmentations\n    image = tf.image.random_flip_left_right(image)\n    image = tf.image.random_hue(image, 0.01)\n    image = tf.image.random_saturation(image, 0.70, 1.30)\n    image = tf.image.random_contrast(image, 0.80, 1.20)\n    image = tf.image.random_brightness(image, 0.10)\n        \n    # ImageNet Normalization\n    image = tf.cast(image, tf.float32)  / 255.0\n    image = (image - IMAGENET_MEAN) / IMAGENET_STD\n    \n    return { 'image': image, 'individual_id': 42 }, { 'individual_id': 42 }","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:37:46.361863Z","iopub.execute_input":"2022-03-20T10:37:46.362200Z","iopub.status.idle":"2022-03-20T10:37:46.371163Z","shell.execute_reply.started":"2022-03-20T10:37:46.362151Z","shell.execute_reply":"2022-03-20T10:37:46.370266Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_train_dataset_augmentation(labels=['individual_id', 'species', 'family']):\n    FNAMES_TRAIN_TFRECORDS = [f'{GCS_DS_PATH}/*.tfrec']\n    \n    dataset = tf.data.TFRecordDataset.from_tensor_slices(np.zeros(32))\n        \n    dataset = dataset.map(decode_tfrecord_demo, num_parallel_calls=AUTO)\n    dataset = dataset.batch(32)\n    \n    return dataset\n\ndataset = get_train_dataset_augmentation()","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:37:46.453875Z","iopub.execute_input":"2022-03-20T10:37:46.454529Z","iopub.status.idle":"2022-03-20T10:37:46.508576Z","shell.execute_reply.started":"2022-03-20T10:37:46.454484Z","shell.execute_reply":"2022-03-20T10:37:46.507788Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"show_batch(get_train_dataset_augmentation())","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:37:46.724660Z","iopub.execute_input":"2022-03-20T10:37:46.725246Z","iopub.status.idle":"2022-03-20T10:37:54.128859Z","shell.execute_reply.started":"2022-03-20T10:37:46.725054Z","shell.execute_reply":"2022-03-20T10:37:54.128147Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dynamic Margins\n\nBased on the [2nd place solution](https://arxiv.org/pdf/2010.05350.pdf) of the Google Landmark Recognition 2021 competition. Margins are computed based on class occurance, making the descriptors of classes with few examples more unique.\n\nThere is a large class inbalance, as shown below, with number of samples per class ranging from 1 to 400.","metadata":{}},{"cell_type":"code","source":"display(train['individual_id'].value_counts().describe().to_frame(name='Value'))","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:43:26.854038Z","iopub.execute_input":"2022-03-20T10:43:26.854572Z","iopub.status.idle":"2022-03-20T10:43:26.882114Z","shell.execute_reply.started":"2022-03-20T10:43:26.854522Z","shell.execute_reply":"2022-03-20T10:43:26.881225Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.figure(figsize=(15,8))\nplt.title('Individual Id Sample Count', size=24)\ntrain['individual_id'].value_counts().value_counts().sort_index().head(25).plot(kind='bar')\nplt.xticks(size=16)\nplt.yticks(size=16)\nplt.xlabel('Frequency', size=18)\nplt.ylabel('Class Size', size=18)\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:55:11.096274Z","iopub.execute_input":"2022-03-20T10:55:11.096596Z","iopub.status.idle":"2022-03-20T10:55:11.433037Z","shell.execute_reply.started":"2022-03-20T10:55:11.096547Z","shell.execute_reply":"2022-03-20T10:55:11.432301Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_dynamic_margins(a=0.40, b=0.05):\n    # Individual Id Value Counts\n    value_counts = train['individual_id'].value_counts().sort_index()\n        \n    # Compute Dynamic Margins\n    dynamic_margins = value_counts.sort_index().values ** -0.25\n    dynamic_margins = a* dynamic_margins + b\n    dynamic_margins = tf.constant(dynamic_margins, dtype=tf.float32)\n    \n    # Sanity Check, class size to margin mapping\n    class_size2margin = pd.Series(data=dynamic_margins, index=value_counts.sort_index())\n    class_size2margin = class_size2margin.drop_duplicates().sort_index()\n    \n    return dynamic_margins, class_size2margin\n\ndynamic_margins, class_size2margin = get_dynamic_margins()","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:43:40.003115Z","iopub.execute_input":"2022-03-20T10:43:40.003438Z","iopub.status.idle":"2022-03-20T10:43:40.058242Z","shell.execute_reply.started":"2022-03-20T10:43:40.003400Z","shell.execute_reply":"2022-03-20T10:43:40.057372Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Dynamic Margins Statistics\ndisplay(pd.Series(dynamic_margins).describe())","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:43:41.173560Z","iopub.execute_input":"2022-03-20T10:43:41.173894Z","iopub.status.idle":"2022-03-20T10:43:41.184331Z","shell.execute_reply.started":"2022-03-20T10:43:41.173855Z","shell.execute_reply":"2022-03-20T10:43:41.183435Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Plot Dynamic Margins Distribution\n# Over Half of the classes have just 1 sample with the maximum margin of 0.45\nplt.figure(figsize=(15, 8))\nplt.title(f'Dynamic Margins Distribution', size=24)\npd.Series(np.flip(np.sort(dynamic_margins))).plot()\nplt.xticks(size=16)\nplt.yticks(size=16)\nplt.xlabel('Class Count', size=18)\nplt.ylabel('Margin', size=18)\nplt.grid()\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:54:23.167990Z","iopub.execute_input":"2022-03-20T10:54:23.168517Z","iopub.status.idle":"2022-03-20T10:54:23.883719Z","shell.execute_reply.started":"2022-03-20T10:54:23.168483Z","shell.execute_reply":"2022-03-20T10:54:23.882857Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Plot Dynamic Margins Distribution\nplt.figure(figsize=(15, 8))\nplt.title(f'Class Size to Margin Mapping', size=24)\nplt.xscale('log')\nclass_size2margin.plot()\nplt.xticks(size=16)\nplt.yticks(size=16)\nplt.xlabel('Class Size', size=18)\nplt.ylabel('Margin', size=18)\nplt.grid()\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:53:33.612003Z","iopub.execute_input":"2022-03-20T10:53:33.612351Z","iopub.status.idle":"2022-03-20T10:53:34.064792Z","shell.execute_reply.started":"2022-03-20T10:53:33.612315Z","shell.execute_reply":"2022-03-20T10:53:34.063945Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Class size to margin mapping\ndisplay(class_size2margin.head(10))","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:44:14.283112Z","iopub.execute_input":"2022-03-20T10:44:14.284109Z","iopub.status.idle":"2022-03-20T10:44:14.290945Z","shell.execute_reply.started":"2022-03-20T10:44:14.284052Z","shell.execute_reply":"2022-03-20T10:44:14.290149Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# ArcMargin Product","metadata":{}},{"cell_type":"code","source":"class ArcMarginPenaltyLogists(tf.keras.layers.Layer):\n    \"\"\"ArcMarginPenaltyLogists\"\"\"\n    def __init__(self, num_classes, dynamic_margins, logist_scale=30, k=1, **kwargs):\n        super(ArcMarginPenaltyLogists, self).__init__(**kwargs)\n        self.num_classes = num_classes\n        if type(dynamic_margins) == float:\n            print(f'Using Static Margin: {dynamic_margins}')\n            self.dynamic_margins = tf.fill(dims=[num_classes], value=dynamic_margins)\n        else:\n            print(f'Using Dynamic Margins, first 5: {dynamic_margins[:5]}')\n            self.dynamic_margins = dynamic_margins\n        self.logist_scale = logist_scale\n        self.k = k\n\n    def build(self, input_shape):\n        initializer_amplitude = 1.0 / tf.math.sqrt(float(self.num_classes))\n        initializer = tf.keras.initializers.random_uniform(-initializer_amplitude, initializer_amplitude)\n        self.w = self.add_variable(\"weights\", shape=[int(input_shape[-1]), self.num_classes * self.k], initializer=initializer)\n\n    def call(self, embds, labels, debug=False):\n        # Dynamic Margins\n        margins = tf.gather(self.dynamic_margins, labels)\n        cos_m = tf.identity(tf.math.cos(margins), name='cos_m')\n        sin_m = tf.identity(tf.math.sin(margins), name='sin_m')\n        th = tf.identity(tf.math.cos(math.pi - margins), name='th')\n        mm = tf.multiply(sin_m, margins, name='mm')\n        if debug:\n            print(f'margins: {margins}')\n            print(f'cos shape: {cos_m.shape}, sin_m shape: {sin_m.shape}, th shape: {th.shape}, mm shape: {mm.shape}')\n        \n        embds = tf.cast(embds, tf.float32)\n        \n        normed_embds = tf.nn.l2_normalize(embds, axis=1, name='normed_embd')\n        normed_w = tf.nn.l2_normalize(self.w, axis=0, name='normed_weights')\n\n        cos_t = tf.matmul(normed_embds, normed_w, name='cos_t')\n        cos_t = tf.reshape(cos_t, shape=[-1, self.num_classes, self.k])\n        cos_t = tf.math.reduce_max(cos_t, axis=2)\n        if debug:\n            print(f'cos_t shape: {cos_t.shape}')\n        sin_t = tf.sqrt(1. - cos_t ** 2, name='sin_t')\n\n        cos_mt = tf.subtract(cos_t * tf.reshape(cos_m, [-1, 1]), sin_t * tf.reshape(sin_m, [-1, 1]), name='cos_mt')\n\n        cos_mt = tf.where(cos_t > tf.reshape(th, [-1, 1]), cos_mt, cos_t - tf.reshape(mm, [-1, 1]))\n\n        mask = tf.one_hot(tf.cast(labels, tf.int32), depth=self.num_classes, name='one_hot_mask')\n\n        logists = tf.where(mask == 1., cos_mt, cos_t)\n        logists = tf.multiply(logists, self.logist_scale, 'arcface_logist')\n\n        return logists","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:44:41.201379Z","iopub.execute_input":"2022-03-20T10:44:41.201703Z","iopub.status.idle":"2022-03-20T10:44:41.221600Z","shell.execute_reply.started":"2022-03-20T10:44:41.201667Z","shell.execute_reply":"2022-03-20T10:44:41.220772Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# DOLG\n\nDOLG: DOLG: Single-Stage Image Retrieval with Deep Orthogonal Fusion of Local and Global Features ([paper](https://arxiv.org/pdf/2108.02927.pdf))\n\nImplementation based on [Pytorch implementation](https://github.com/ChristofHenkel/kaggle-landmark-2021-1st-place/blob/main/models/ch_mdl_dolg_efficientnet.py) published by Christof Henkel on GitHub.","metadata":{}},{"cell_type":"code","source":"class GeM(tf.keras.layers.Layer):\n    def __init__(self, init_norm=3.0, **kwargs):\n        super(GeM, self).__init__(**kwargs)\n        self.init_norm = init_norm\n        self.gap2d = tf.keras.layers.GlobalAveragePooling2D()\n\n    def build(self, input_shape):\n        super(GeM, self).build(input_shape)\n        feature_size = input_shape[-1]\n        self.p = self.add_weight(\n                name = \"norms\",\n                shape = feature_size,\n                initializer = tf.keras.initializers.constant(self.init_norm),\n                trainable = True,\n            )\n\n    def call(self, inputs):\n        x = tf.math.maximum(inputs, EPS)\n        x = tf.pow(x, self.p)\n\n        x = self.gap2d(x)\n        x = tf.pow(x, 1.0 / self.p)\n\n        return x","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:45:58.785066Z","iopub.execute_input":"2022-03-20T10:45:58.785413Z","iopub.status.idle":"2022-03-20T10:45:58.794705Z","shell.execute_reply.started":"2022-03-20T10:45:58.785376Z","shell.execute_reply":"2022-03-20T10:45:58.794025Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Multi-Atrous Branch\nclass MultiAtrous(tf.keras.layers.Layer):\n    def __init__(self, dolg_s, upsampling=1, kernel_size=3, padding=\"same\",  **kwargs):\n        super(MultiAtrous, self).__init__(**kwargs)\n        self.d0 = tf.keras.layers.Conv2D(dolg_s // 2, 3, dilation_rate=(3,3), padding='same')\n        self.d1 = tf.keras.layers.Conv2D(dolg_s // 2, 3, dilation_rate=(6,6), padding='same')\n        self.d2 = tf.keras.layers.Conv2D(dolg_s // 2, 3, dilation_rate=(9,9), padding='same')\n        self.conv1 = tf.keras.layers.Conv2D(dolg_s, kernel_size=1)\n        \n    @tf.function()\n    def call(self, inputs, training=None, **kwargs):\n        x0 = self.d0(inputs)\n        x1 = self.d1(inputs)\n        x2 = self.d2(inputs)\n        x = tf.keras.layers.Concatenate(axis=3)([x0,x1,x2])\n        x = self.conv1(x)\n        x = tf.keras.activations.relu(x)\n        return x\n            \n    def get_config(self):\n        config = {\n            'dilation_rates': self.dilation_rates,\n            'kernel_size'   : self.kernel_size,\n            'padding'       : self.padding,\n            'upsampling'    : self.upsampling\n        }\n        base_config = super(MultiAtrous, self).get_config()\n        return dict(list(base_config.items()) + list(config.items()))","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:45:59.434336Z","iopub.execute_input":"2022-03-20T10:45:59.434634Z","iopub.status.idle":"2022-03-20T10:45:59.446872Z","shell.execute_reply.started":"2022-03-20T10:45:59.434605Z","shell.execute_reply":"2022-03-20T10:45:59.446247Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class SpatialAttention2d(tf.keras.layers.Layer):\n    def __init__(self, dolg_s, **kwargs):\n        super(SpatialAttention2d, self).__init__(**kwargs)\n        self.conv1 = tf.keras.layers.Conv2D(dolg_s, 1)\n        self.bn = tf.keras.layers.BatchNormalization()\n        self.conv2 = tf.keras.layers.Conv2D(1, 1)\n\n    @tf.function()\n    def call(self, x):\n        x = self.conv1(x)\n        x = self.bn(x)\n        \n        feature_map_norm, _ = tf.linalg.normalize(x, ord=2, axis=3)\n        \n        x = tf.keras.activations.relu(x)\n        x = self.conv2(x)\n        \n        att_score = tf.keras.activations.softplus(x)\n        \n        x = att_score * feature_map_norm\n\n        return x","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:45:59.929861Z","iopub.execute_input":"2022-03-20T10:45:59.930694Z","iopub.status.idle":"2022-03-20T10:45:59.940015Z","shell.execute_reply.started":"2022-03-20T10:45:59.930647Z","shell.execute_reply":"2022-03-20T10:45:59.939269Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class OrthogonalFusion(tf.keras.layers.Layer):\n    def __init__(self, **kwargs):\n        super().__init__(**kwargs)\n        \n    @tf.function()\n    def call(self, inputs):\n        fl, fg = inputs\n        fl = tf.transpose(fl, [0,3,1,2])\n        \n        bs, c, w, h = fl.shape\n        \n        fl_b = tf.reshape(fl, [tf.shape(fl)[0],c,w*h])\n        fl_dot_fg = tf.matmul(fg[:,tf.newaxis,:] ,fl_b)\n       \n        fl_dot_fg = tf.reshape(fl_dot_fg, [tf.shape(fl_dot_fg)[0],1,w,h])\n        \n        fg_norm = tf.norm(fg, ord=2, axis=1)\n        \n        fl_proj = (fl_dot_fg / fg_norm[:,tf.newaxis,tf.newaxis,tf.newaxis]) * fg[:,:,tf.newaxis,tf.newaxis]\n        fl_orth = fl - fl_proj\n        \n        fg_rep = tf.tile(fg[:,:,tf.newaxis,tf.newaxis], multiples=(1,1,w,h))\n        f_fused = tf.keras.layers.Concatenate(axis=1)([fl_orth, fg_rep])\n        \n        # Transpose\n        f_fused = tf.transpose(f_fused, [0,2,3,1])\n        \n        return f_fused","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:46:00.103024Z","iopub.execute_input":"2022-03-20T10:46:00.103488Z","iopub.status.idle":"2022-03-20T10:46:00.114940Z","shell.execute_reply.started":"2022-03-20T10:46:00.103456Z","shell.execute_reply":"2022-03-20T10:46:00.114095Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class GlobalBranch(tf.keras.layers.Layer):\n    def __init__(self, dolg_s, **kwargs):\n        super().__init__(**kwargs)\n        self.conv2d = tf.keras.layers.Conv2D(dolg_s, 1, name='global_conv2d')\n        self.bn = tf.keras.layers.BatchNormalization()\n        self.pool = tf.keras.layers.GlobalAveragePooling2D()\n        \n    @tf.function()\n    def call(self, inputs):\n        x = self.conv2d(inputs)\n        x = self.bn(x)\n        x = tf.nn.silu(x)\n        x = self.pool(x)\n        \n        return x","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:46:00.563674Z","iopub.execute_input":"2022-03-20T10:46:00.564109Z","iopub.status.idle":"2022-03-20T10:46:00.570896Z","shell.execute_reply.started":"2022-03-20T10:46:00.564076Z","shell.execute_reply":"2022-03-20T10:46:00.570237Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class DolgBranch(tf.keras.layers.Layer):\n    def __init__(self, dolg_s, idx, **kwargs):\n        super().__init__(name=f'dolg_branch_{idx}', **kwargs)\n        dolg_s = int(dolg_s)\n        # Local\n        self.mam = MultiAtrous(dolg_s, name=f'mam_{idx}')\n        self.sa2d = SpatialAttention2d(dolg_s, name=f'sa2d_{idx}')\n        # Global\n        self.global_branch = GlobalBranch(dolg_s, name=f'g_{idx}')\n        # Orthogonal Fusion\n        self.orthogonal_fusion = OrthogonalFusion()\n        # Pooling\n        self.pool = tf.keras.layers.GlobalAveragePooling2D()\n        \n    @tf.function()\n    def call(self, inputs):\n        inputs_l, inputs_g = inputs\n        # Local\n        l = self.mam(inputs_l)\n        l = self.sa2d(l)\n        # Global\n        g = self.global_branch(inputs_g)\n        # Orthogonal Fusion\n        f = self.orthogonal_fusion([l, g])\n        # Pooling\n        descriptor = self.pool(f)\n        \n        return descriptor","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:46:00.906130Z","iopub.execute_input":"2022-03-20T10:46:00.906997Z","iopub.status.idle":"2022-03-20T10:46:00.916959Z","shell.execute_reply.started":"2022-03-20T10:46:00.906939Z","shell.execute_reply":"2022-03-20T10:46:00.915953Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model\n\nEfficientNetV2 models published by Google on [GitHub](https://github.com/google/automl/tree/master/efficientnetv2)","metadata":{}},{"cell_type":"code","source":"GCS_WEIGHTS_PATH = KaggleDatasets().get_gcs_path('efficientnetv2-pretrained-imagenet21k-weights')","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:46:41.130666Z","iopub.execute_input":"2022-03-20T10:46:41.130969Z","iopub.status.idle":"2022-03-20T10:46:41.428534Z","shell.execute_reply.started":"2022-03-20T10:46:41.130938Z","shell.execute_reply":"2022-03-20T10:46:41.427491Z"},"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    with strategy.scope():\n        # Input\n        image = tf.keras.layers.Input(INPUT_SHAPE, name='image', dtype=tf.float32)\n        individual_id = tf.keras.layers.Input([], name='individual_id_input', dtype=tf.int32)\n        \n        \n        # EfficientNetV2 CNN\n        cnn = effnetv2_model.get_model(f'efficientnetv2-{EFNV2_SIZE}', include_top=False, weights=None)\n        \n        # Load Pretrained ImageNet21K Finetuned Imagenet1K Weights\n        if TPU:\n            WEIGHT_PATH = f'{GCS_WEIGHTS_PATH}/efficientnetv2-{EFNV2_SIZE}-21k-ft1k'\n            ckpt = tf.train.latest_checkpoint(WEIGHT_PATH)\n            cnn.load_weights(ckpt)\n        \n        # CNN Outputs\n        embedding, fm5, fm4, fm3, fm2, fm1 = cnn(image, with_endpoints=True)\n        print(f'embedding: {embedding.shape}, fm5: {fm5.shape}, fm4: {fm4.shape}, fm3: {fm3.shape}, fm2: {fm2.shape}, fm1: {fm1.shape}')\n        \n        # DOLG Branches\n        descriptor = DolgBranch(DOLG_SIZE, 1)([fm2, fm1])\n        \n        # Dense Layer\n        descriptor = tf.keras.layers.Dropout(0.00)(descriptor)\n        descriptor = tf.keras.layers.Dense(EMBEDDING_SIZE, name='descriptor_dense')(descriptor)\n        \n        # ArcMarginProduct\n        outputs = ArcMarginPenaltyLogists(N_INDIVIDUAL_IDS, dynamic_margins=dynamic_margins, k=1, name='individual_id')(descriptor, individual_id)\n        \n        model = tf.keras.models.Model(inputs=[image, individual_id], outputs=outputs)\n        \n        # OPTIMIZER\n        optimizer = tf.optimizers.Adam()\n        \n        # LOSS\n        loss = {\n            'individual_id': tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n        }\n        # METRICS\n        metrics =[\n            tf.keras.metrics.SparseTopKCategoricalAccuracy(k=5, name='top5acc'),\n            tf.keras.metrics.SparseTopKCategoricalAccuracy(k=1, name='top1acc'),\n        ]\n\n        # Compile Model\n        model.compile(optimizer=optimizer, loss=loss, metrics=metrics)\n\n        return model","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:46:41.971255Z","iopub.execute_input":"2022-03-20T10:46:41.972184Z","iopub.status.idle":"2022-03-20T10:46:41.985959Z","shell.execute_reply.started":"2022-03-20T10:46:41.972133Z","shell.execute_reply":"2022-03-20T10:46:41.985219Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = get_model()","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:46:52.294680Z","iopub.execute_input":"2022-03-20T10:46:52.295264Z","iopub.status.idle":"2022-03-20T10:47:44.773263Z","shell.execute_reply.started":"2022-03-20T10:46:52.295224Z","shell.execute_reply":"2022-03-20T10:47:44.772300Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# The EfficientNetV2-XL Model is Huge with over 200 million parameters!\nmodel.summary()","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:47:44.774878Z","iopub.execute_input":"2022-03-20T10:47:44.775134Z","iopub.status.idle":"2022-03-20T10:47:44.890958Z","shell.execute_reply.started":"2022-03-20T10:47:44.775104Z","shell.execute_reply":"2022-03-20T10:47:44.890156Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"tf.keras.utils.plot_model(model, show_shapes=True, show_dtype=False, show_layer_names=True, expand_nested=False)","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:47:44.892190Z","iopub.execute_input":"2022-03-20T10:47:44.892404Z","iopub.status.idle":"2022-03-20T10:47:45.923893Z","shell.execute_reply.started":"2022-03-20T10:47:44.892378Z","shell.execute_reply":"2022-03-20T10:47:45.922948Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Learning Rate Scheduler","metadata":{}},{"cell_type":"code","source":"TRAIN_STEPS_PER_EPOCH = N_SAMPLES // BATCH_SIZE\nprint(f'N_EPOCHS: {N_EPOCHS}, TRAIN_STEPS_PER_EPOCH: {TRAIN_STEPS_PER_EPOCH}')","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:47:45.926260Z","iopub.execute_input":"2022-03-20T10:47:45.926520Z","iopub.status.idle":"2022-03-20T10:47:45.931605Z","shell.execute_reply.started":"2022-03-20T10:47:45.926489Z","shell.execute_reply":"2022-03-20T10:47:45.930793Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def lrfn(current_step, num_warmup_steps, lr_max, num_cycles=0.50, num_training_steps=N_EPOCHS):\n    \n    if current_step < num_warmup_steps:\n        return lr_max * 0.5 ** (num_warmup_steps - current_step)\n    else:\n        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))\n\n        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress))) * lr_max","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:47:45.932695Z","iopub.execute_input":"2022-03-20T10:47:45.932923Z","iopub.status.idle":"2022-03-20T10:47:45.948597Z","shell.execute_reply.started":"2022-03-20T10:47:45.932894Z","shell.execute_reply":"2022-03-20T10:47:45.947510Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def plot_lr_schedule(lr_schedule, name):\n    fig = plt.figure(figsize=(20,8))\n    plt.plot([None] + lr_schedule + [None])\n    # X Labels\n    x = np.arange(N_EPOCHS + 2)\n    x_axis_labels = [None] + list(map(str, np.arange(1, N_EPOCHS+1))) + [None]\n    plt.xlim([0, N_EPOCHS + 1])\n    plt.xticks(x, x_axis_labels, size=12) # set tick step to 1 and let x axis start at 1\n    plt.yticks(size=12)\n    \n    # Increase y-limit for better readability\n    plt.ylim([0, max(lr_schedule) * 1.1])\n    \n    # Title\n    schedule_info = f'start: {lr_schedule[0]:.1E}, max: {max(lr_schedule):.1E}, final: {lr_schedule[-1]:.1E}'\n    plt.title(f'Step Learning Rate Schedule {name}, {schedule_info}', size=18, pad=12)\n    \n    # Plot Learning Rates\n    for x, val in enumerate(lr_schedule):\n        if x < len(lr_schedule) - 1:\n            if lr_schedule[x - 1] < val:\n                ha = 'right'\n            else:\n                ha = 'left'\n        elif x == 0:\n            ha = 'right'\n        else:\n            ha = 'left'\n        plt.plot(x + 1, val, 'o', color='black');\n        offset_y = (max(lr_schedule) - min(lr_schedule)) * 0.02\n        plt.annotate(f'{val:.1E}', xy=(x + 1, val + offset_y), size=12, ha=ha)\n    \n    plt.xlabel('Epoch', size=16, labelpad=5)\n    plt.ylabel('Learning Rate', size=16, labelpad=5)\n    plt.grid()\n    plt.show()\n\n# Learning rate for encoder\nLR_SCHEDULE = [lrfn(step, num_warmup_steps=3, lr_max=5e-4, num_cycles=0.50) for step in range(N_EPOCHS)]\nplot_lr_schedule(LR_SCHEDULE, 'Model')","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:56:07.789797Z","iopub.execute_input":"2022-03-20T10:56:07.790894Z","iopub.status.idle":"2022-03-20T10:56:08.215448Z","shell.execute_reply.started":"2022-03-20T10:56:07.790837Z","shell.execute_reply":"2022-03-20T10:56:08.214358Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Callbacks","metadata":{}},{"cell_type":"code","source":"# Make Checkpoints for Models based on best training loss\nmodel_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(\n    'model_best.h5', monitor='loss', verbose=1, save_best_only=True, save_weights_only=True\n)\nmodel_checkpoint_callback.set_model(model)\n\n# Learning Rate Scheduler\nlearning_rate_callback = tf.keras.callbacks.LearningRateScheduler(lambda epoch: LR_SCHEDULE[epoch], verbose=1)","metadata":{"execution":{"iopub.status.busy":"2022-03-16T14:30:02.866829Z","iopub.execute_input":"2022-03-16T14:30:02.867089Z","iopub.status.idle":"2022-03-16T14:30:02.87312Z","shell.execute_reply.started":"2022-03-16T14:30:02.86706Z","shell.execute_reply":"2022-03-16T14:30:02.872041Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Training\nhistory = model.fit(\n    get_train_dataset(),\n    steps_per_epoch = TRAIN_STEPS_PER_EPOCH,\n    epochs = N_EPOCHS,\n    verbose = 1,\n    callbacks = [\n        learning_rate_callback,\n        model_checkpoint_callback,\n    ],\n)","metadata":{"execution":{"iopub.status.busy":"2022-03-16T14:30:02.874331Z","iopub.execute_input":"2022-03-16T14:30:02.874581Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Load Best Weights\nmodel.load_weights('model_best.h5')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training History","metadata":{}},{"cell_type":"code","source":"def plot_history_metric(metric, f_best=np.argmax, yscale='linear'):\n    x = np.arange(1, len(history.history[metric]) + 1)\n    y_train = history.history[metric]\n    plt.figure(figsize=(20, 8))\n    # TRAIN\n    plt.plot(x, y_train, color='tab:blue', lw=3, label='train')\n    plt.title(f'Training {metric}', fontsize=24, pad=10)\n    plt.ylabel(metric, fontsize=20, labelpad=10)\n    plt.xlabel('epoch', fontsize=20, labelpad=10)\n    plt.xticks([1] + np.arange(5, N_EPOCHS + 1, 5).tolist(), fontsize=16) # set tick step to 1 and let x axis start at 1\n    plt.yticks(fontsize=16)\n    plt.yscale(yscale)\n    \n    # Train Best Marker\n    x_best = f_best(y_train)\n    y_best = y_train[x_best]\n    plt.scatter(x_best + 1, y_best, color='purple', s=100, marker='o', label=f'train best: {y_best:.4f}')\n \n    if f'val_{metric}' in history.history:\n        y_val = history.history[f'val_{metric}']\n       # Validation Best Marker\n        plt.plot(x, y_val, color='tab:orange', lw=3, label='validation')\n        # VALIDATION\n        x_best = f_best(y_val)\n        y_best = y_val[x_best]\n        plt.scatter(x_best + 1, y_best, color='red', s=100, marker='o', label=f'validation best: {y_best:.4f}')\n    \n    plt.grid()\n    plt.legend(prop={'size': 18})\n    plt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_history_metric('loss', f_best=np.argmin)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_history_metric('top5acc', f_best=np.argmax)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plot_history_metric('top1acc', f_best=np.argmax)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Create Embeddings Model\n\nCreates the embedding model where the output is not the classifier, but the descriptor.","metadata":{}},{"cell_type":"code","source":"# Show Model Names\nfor idx, l in enumerate(model.layers):\n    print(f'{idx} | \\t{l.name}')","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:47:58.357543Z","iopub.execute_input":"2022-03-20T10:47:58.358016Z","iopub.status.idle":"2022-03-20T10:47:58.364631Z","shell.execute_reply.started":"2022-03-20T10:47:58.357967Z","shell.execute_reply":"2022-03-20T10:47:58.363665Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with strategy.scope():\n    # Input\n    image = tf.keras.layers.Input(INPUT_SHAPE, name='image', dtype=tf.float32)\n\n    # EfficientNet\n    embedding, fm5, fm4, fm3, fm2, fm1 = model.layers[1](image, with_endpoints=True)\n    \n    # DOLG Branches\n    descriptor = model.layers[2]([fm2, fm1])\n    \n    # Descriptor\n    outputs = model.layers[4](descriptor)\n    \n    model_embedding = tf.keras.Model(inputs=image, outputs=outputs)\n    model_embedding.trainable = False","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:47:58.629109Z","iopub.execute_input":"2022-03-20T10:47:58.629917Z","iopub.status.idle":"2022-03-20T10:48:18.499132Z","shell.execute_reply.started":"2022-03-20T10:47:58.629881Z","shell.execute_reply":"2022-03-20T10:48:18.498317Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model_embedding.summary()","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:48:18.500674Z","iopub.execute_input":"2022-03-20T10:48:18.500917Z","iopub.status.idle":"2022-03-20T10:48:18.602959Z","shell.execute_reply.started":"2022-03-20T10:48:18.500887Z","shell.execute_reply":"2022-03-20T10:48:18.602045Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# The embedding model outputs the descriptor of size 2048\ntf.keras.utils.plot_model(model_embedding, show_shapes=True, show_dtype=True, show_layer_names=True, expand_nested=False)","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:49:00.217951Z","iopub.execute_input":"2022-03-20T10:49:00.218911Z","iopub.status.idle":"2022-03-20T10:49:00.406804Z","shell.execute_reply.started":"2022-03-20T10:49:00.218850Z","shell.execute_reply":"2022-03-20T10:49:00.405705Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model_embedding.save_weights('model_embedding.h5')","metadata":{"execution":{"iopub.status.busy":"2022-03-20T10:48:18.778317Z","iopub.execute_input":"2022-03-20T10:48:18.778575Z","iopub.status.idle":"2022-03-20T10:48:21.544848Z","shell.execute_reply.started":"2022-03-20T10:48:18.778542Z","shell.execute_reply":"2022-03-20T10:48:21.544091Z"},"trusted":true},"execution_count":null,"outputs":[]}]}