{"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":"code","source":"import os\nIS_COLAB = not os.path.exists('/kaggle/input')\nprint(IS_COLAB) ","metadata":{"id":"p9yESN6Ng8Wt","outputId":"81b83105-3043-4313-edc2-cbed72b15450","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install efficientnet tensorflow_addons > /dev/null","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import tensorflow as tf\ntry:\n    # TPU detection. No parameters necessary if TPU_NAME environment variable is\n    # set: this is always the case on Kaggle.\n    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n    print('Running on TPU ', tpu.master())\nexcept ValueError:\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.TPUStrategy(tpu)\nelse:\n    # Default distribution strategy in Tensorflow. Works on CPU and single GPU.\n    strategy = tf.distribute.get_strategy()\n\nAUTO = tf.data.experimental.AUTOTUNE\nprint(\"REPLICAS: \", strategy.num_replicas_in_sync)","metadata":{"id":"d59NPeZ0u-_B","outputId":"a354880e-e82e-4461-acf4-882662760afb","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# tpu_worker = os.environ['COLAB_TPU_ADDR'].replace('8470', '8466') \n# print(tf.profiler.experimental.client.monitor(tpu_worker,1))","metadata":{"id":"_EL_aDUzhKzn","outputId":"2c6a41e1-59c6-4bf0-d2b3-200786b8b282","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if IS_COLAB:\n  from google.colab import drive\n  drive.mount('/content/drive')\nelse:\n  from kaggle_datasets import KaggleDatasets","metadata":{"id":"jZBp3Th1i4c6","outputId":"e56cc635-dff2-4c15-ee82-b58d2037d358","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install -q efficientnet\n!pip install tensorflow_addons\nimport re\nimport os\nimport numpy as np\nimport pandas as pd\nimport random\nimport math\nimport tensorflow as tf\nimport efficientnet.tfkeras as efn\nfrom sklearn import metrics\nfrom sklearn.model_selection import KFold, train_test_split\nfrom tensorflow.keras import backend as K\nimport tensorflow_addons as tfa\nfrom tqdm.auto import tqdm\nimport matplotlib.pyplot as plt\nimport pickle\nimport json\nimport tensorflow_hub as tfhub\nfrom datetime import datetime\nEXPERIMENT = 4","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","id":"0RDuS-Zyu---","outputId":"0019dc04-91bf-44c1-ff69-bc18337ca64c","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if IS_COLAB:\n    ## Clean Folder if empty\n    clean=False\n    if clean:\n      for folder in os.listdir('/content/drive/MyDrive/Kaggle/GLR-2021/experiments-1/'):\n        if len(os.listdir(f'/content/drive/MyDrive/Kaggle/GLR-2021/experiments-1/{folder}'))==0:\n          print(\"cleaning\",folder)\n          !rm -r '/content/drive/MyDrive/Kaggle/GLR-2021/experiments-1/{folder}'","metadata":{"id":"G4LZ_WKXkzxl","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"save_dir = '.'\nif IS_COLAB:\n    run_ts = datetime.now().strftime('%Y%m%d-%H%M%S')\n    save_dir = f'/content/drive/MyDrive/Kaggle/GLR-2021/experiments-{EXPERIMENT}/{run_ts}'\n    print(run_ts)\n    !mkdir -p {save_dir}","metadata":{"id":"S2MC5VAhgdz1","outputId":"eb488715-acd8-4401-c2cc-4c9c12145b93","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DEBUG = False # True","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class config:\n    \n    \n    SEED = 42\n    FOLD_TO_RUN = 0\n    FOLDS = 5\n    DEBUG = DEBUG # raw False\n    EVALUATE = True\n    RESUME = False# True\n    RESUME_EPOCH = 39\n    # model_path = f'/content/drive/MyDrive/Kaggle/GLR-2021/experiments-{EXPERIMENT}/20210831-200946/'\n    model_path = f'../input/glr-eff-v2-m-arcface-retraining-at-640/'\n    \n    ### Dataset\n    # dataset = 'v2'  # one of 'v2', 'v2c', 'comp'\n    dataset = 'comp'\n    BATCH_SIZE = 32 * strategy.num_replicas_in_sync\n    IMAGE_SIZE = 256 if DEBUG else 640\n    \n    ### Model\n    model_type = 'effnetv2'  # One of effnetv1, effnetv2\n    EFF_NET = 0 if DEBUG else 6\n    EFF_NETV2 = 'm-21k-ft1k'\n    FREEZE_BATCH_NORM = False\n    head = 'arcface' # one of arcface, curricular-face\n    EPOCHS = 1 if DEBUG else 1\n    LR = 0.001\n    message='retraining 640 epoch 2'\n    \n    ### Augmentations\n    PRECROP_IMAGE_SIZE = 256 if DEBUG else 512\n    ROT_ = 10.0\n    SHR_ = 2.0\n    HZOOM_ = 8.0\n    WZOOM_ = 8.0\n    HSHIFT_ = 8.0\n    WSHIFT_ = 8.0\n    CUTOUT = False\n    save_dir = save_dir","metadata":{"id":"Vky-cNilgdz2","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if config.model_type == 'effnetv2' and (not IS_COLAB):\n    EFFNETV2_ROOT = KaggleDatasets().get_gcs_path('efficientnetv2-tfhub-weight-files')\n    print(EFFNETV2_ROOT)\nelse:\n    ## Run this cell on kaggle to get updated gcs path\n    EFFNETV2_ROOT = 'gs://kds-9b67f7a880034be953f1ea1d0ca5dbd5b5fc5c95bc6337d46401f403'","metadata":{"id":"xmbduQv_gdz2","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if config.EPOCHS>3:\n    ## Save Last 3 epochs\n    config.SNAPSHOT_THRESHOLD = config.EPOCHS-3\nelse:\n    config.SNAPSHOT_THRESHOLD = 0","metadata":{"id":"WFGCZgXagdz2","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"config.SNAPSHOT_THRESHOLD = 99","metadata":{"id":"WlX6psBH1YbD","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if config.dataset=='comp':\n    config.N_CLASSES = 81313\nelif config.dataset=='v2':\n    config.N_CLASSES = 203094","metadata":{"id":"dkAdsz1Ugdz2","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"with open(config.save_dir+'/config.json', 'w') as fp:\n    json.dump({x:dict(config.__dict__)[x] for x in dict(config.__dict__) if not x.startswith('_')}, fp)","metadata":{"id":"XOfbho-ugdz3","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Dataset","metadata":{"id":"5Q3Z9wHzgdz3"}},{"cell_type":"code","source":"# FOLDS = 5\n# FOLD_TO_RUN = 0\n# range(FOLDS)[FOLD_TO_RUN-1]","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if config.dataset=='v2':\n    TRAINING_FILENAMES = [f'gs://landmark-2021/train/landmark-2021-gld2-{x}.tfrec' for x in range(20) if x!=config.FOLD_TO_RUN]\n    VALIDATION_FILENAMES = [f'gs://landmark-2021/train/landmark-2021-gld2-{config.FOLD_TO_RUN}.tfrec']\nelif config.dataset=='comp':\n#     GCS_PATHS = {\n#         0: '/kaggle/input/landmark-recognition-2021-tfrecords-fold0',\n#         1: '/kaggle/input/landmark-recognition-2021-tfrecords-fold1',\n#         2: '/kaggle/input/landmark-recognition-2021-tfrecords-fold2',\n#         3: '/kaggle/input/landmark-recognition-2021-tfrecords-fold3',\n#         4: '/kaggle/input/landmark-recognition-2021-tfrecords-fold4'\n#         }\n    GCS_PATHS = {}\n    for 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        GCS_PATHS[i] =  gcs_path \n    \n    TRAINING_FILENAMES = []\n    VALIDATION_FILENAMES = []\n    \n#     raw --------------------\n    for fold in range(5):\n        GCS_PATH = GCS_PATHS[fold]\n        if fold==config.FOLD_TO_RUN:\n#             VALIDATION_FILENAMES += tf.io.gfile.glob(gcs_path + '/*.tfrec')\n#         TRAINING_FILENAMES += tf.io.gfile.glob(gcs_path + '/*.tfrec')\n            VALIDATION_FILENAMES += tf.io.gfile.glob(GCS_PATH + '/*.tfrec')\n        TRAINING_FILENAMES += tf.io.gfile.glob(GCS_PATH + '/*.tfrec')\n        \n#     for fold in range(config.FOLDS):\n#         GCS_PATH = GCS_PATHS[fold]\n#         if fold==range(config.FOLDS)[config.FOLD_TO_RUN-1]: # 往前一个数据作为验证集\n#             VALIDATION_FILENAMES += tf.io.gfile.glob(gcs_path + '/*.tfrec')\n#         TRAINING_FILENAMES += tf.io.gfile.glob(gcs_path + '/*.tfrec')\n            \nif config.DEBUG:\n    TRAINING_FILENAMES = [TRAINING_FILENAMES[0]]\n    VALIDATION_FILENAMES = [VALIDATION_FILENAMES[0]]\n    \nprint(len(TRAINING_FILENAMES))\nprint(len(VALIDATION_FILENAMES))","metadata":{"id":"6wyHpC5Fgdz4","outputId":"5e55bf7f-4e1c-46aa-8273-ab8ddafa90a0","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"GCS_PATHS","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_mat(rotation, shear, height_zoom, width_zoom, height_shift, width_shift):\n    # returns 3x3 transformmatrix which transforms indicies\n        \n    # CONVERT DEGREES TO RADIANS\n    rotation = math.pi * rotation / 180.\n    shear    = math.pi * shear    / 180.\n\n    def get_3x3_mat(lst):\n        return tf.reshape(tf.concat([lst],axis=0), [3,3])\n    \n    # ROTATION MATRIX\n    c1   = tf.math.cos(rotation)\n    s1   = tf.math.sin(rotation)\n    one  = tf.constant([1],dtype='float32')\n    zero = tf.constant([0],dtype='float32')\n    \n    rotation_matrix = get_3x3_mat([c1,   s1,   zero, \n                                   -s1,  c1,   zero, \n                                   zero, zero, one])    \n    # SHEAR MATRIX\n    c2 = tf.math.cos(shear)\n    s2 = tf.math.sin(shear)    \n    \n    shear_matrix = get_3x3_mat([one,  s2,   zero, \n                                zero, c2,   zero, \n                                zero, zero, one])        \n    # ZOOM MATRIX\n    zoom_matrix = get_3x3_mat([one/height_zoom, zero,           zero, \n                               zero,            one/width_zoom, zero, \n                               zero,            zero,           one])    \n    # SHIFT MATRIX\n    shift_matrix = get_3x3_mat([one,  zero, height_shift, \n                                zero, one,  width_shift, \n                                zero, zero, one])\n    \n    return K.dot(K.dot(rotation_matrix, shear_matrix), \n                 K.dot(zoom_matrix,     shift_matrix))\n\n\ndef transform(image, DIM=256):    \n    # input image - is one image of size [dim,dim,3] not a batch of [b,dim,dim,3]\n    # output - image randomly rotated, sheared, zoomed, and shifted\n    XDIM = DIM%2 #fix for size 331\n    \n    rot = config.ROT_ * tf.random.normal([1], dtype='float32')\n    shr = config.SHR_ * tf.random.normal([1], dtype='float32') \n    h_zoom = 1.0 + tf.random.normal([1], dtype='float32') / config.HZOOM_\n    w_zoom = 1.0 + tf.random.normal([1], dtype='float32') / config.WZOOM_\n    h_shift = config.HSHIFT_ * tf.random.normal([1], dtype='float32') \n    w_shift = config.WSHIFT_ * tf.random.normal([1], dtype='float32') \n\n    # GET TRANSFORMATION MATRIX\n    m = get_mat(rot,shr,h_zoom,w_zoom,h_shift,w_shift) \n\n    # LIST DESTINATION PIXEL INDICES\n    x   = tf.repeat(tf.range(DIM//2, -DIM//2,-1), DIM)\n    y   = tf.tile(tf.range(-DIM//2, DIM//2), [DIM])\n    z   = tf.ones([DIM*DIM], dtype='int32')\n    idx = tf.stack( [x,y,z] )\n    \n    # ROTATE DESTINATION PIXELS ONTO ORIGIN PIXELS\n    idx2 = K.dot(m, tf.cast(idx, dtype='float32'))\n    idx2 = K.cast(idx2, dtype='int32')\n    idx2 = K.clip(idx2, -DIM//2+XDIM+1, DIM//2)\n    \n    # FIND ORIGIN PIXEL VALUES           \n    idx3 = tf.stack([DIM//2-idx2[0,], DIM//2-1+idx2[1,]])\n    d    = tf.gather_nd(image, tf.transpose(idx3))\n        \n    return tf.reshape(d,[DIM, DIM,3])","metadata":{"id":"rV979faMgdz4","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Function to get our f1 score\ndef f1_score(y_true, y_pred):\n    y_true = y_true.apply(lambda x: set(x.split()))\n    y_pred = y_pred.apply(lambda x: set(x.split()))\n    intersection = np.array([len(x[0] & x[1]) for x in zip(y_true, y_pred)])\n    len_y_pred = y_pred.apply(lambda x: len(x)).values\n    len_y_true = y_true.apply(lambda x: len(x)).values\n    f1 = 2 * intersection / (len_y_pred + len_y_true)\n    return f1\n\n# Function to seed everything\ndef seed_everything(seed):\n    random.seed(seed)\n    np.random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    tf.random.set_seed(seed)\n    \ndef arcface_format(posting_id, image, label_group, matches):\n    return posting_id, {'inp1': image, 'inp2': label_group}, label_group, matches\n\n# Data augmentation function\ndef data_augment(posting_id, image, label_group, matches):\n\n    ### CUTOUT\n    if tf.random.uniform([])>0.5 and config.CUTOUT:\n      N_CUTOUT = 6\n      for cutouts in range(N_CUTOUT):\n        if tf.random.uniform([])>0.5:\n           DIM = config.IMAGE_SIZE\n           CUTOUT_LENGTH = DIM//8\n           x1 = tf.cast( tf.random.uniform([],0,DIM-CUTOUT_LENGTH),tf.int32)\n           x2 = tf.cast( tf.random.uniform([],0,DIM-CUTOUT_LENGTH),tf.int32)\n           filter_ = tf.concat([tf.zeros((x1,CUTOUT_LENGTH)),tf.ones((CUTOUT_LENGTH,CUTOUT_LENGTH)),tf.zeros((DIM-x1-CUTOUT_LENGTH,CUTOUT_LENGTH))],axis=0)\n           filter_ = tf.concat([tf.zeros((DIM,x2)),filter_,tf.zeros((DIM,DIM-x2-CUTOUT_LENGTH))],axis=1)\n           cutout = tf.reshape(1-filter_,(DIM,DIM,1))\n           image = cutout*image\n\n    image = tf.image.random_flip_left_right(image)\n    # image = tf.image.random_flip_up_down(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    return posting_id, image, label_group, matches\n\n# Function to decode our images\ndef decode_image(image_data):\n    image = tf.image.decode_jpeg(image_data, channels = 3)\n    image = tf.image.resize(image, [config.IMAGE_SIZE,config.IMAGE_SIZE])\n    image = tf.cast(image, tf.float32) / 255.0\n    return image\n\n# This function parse our images and also get the target variable\ndef read_labeled_tfrecord(example):\n    LABELED_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#         \"matches\": tf.io.FixedLenFeature([], tf.string)\n    }\n\n    example = tf.io.parse_single_example(example, LABELED_TFREC_FORMAT)\n    posting_id = example['image_name']\n    image = decode_image(example['image'])\n#     label_group = tf.one_hot(tf.cast(example['label_group'], tf.int32), depth = N_CLASSES)\n    label_group = tf.cast(example['target'], tf.int32)\n#     matches = example['matches']\n    matches = 1\n    return posting_id, image, label_group, matches\n\n# This function loads TF Records and parse them into tensors\ndef load_dataset(filenames, ordered = False):\n    \n    ignore_order = tf.data.Options()\n    if not ordered:\n        ignore_order.experimental_deterministic = False \n        \n    dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads = AUTO)\n#     dataset = dataset.cache()\n    dataset = dataset.with_options(ignore_order)\n    dataset = dataset.map(read_labeled_tfrecord, num_parallel_calls = AUTO) \n    return dataset\n\n# This function is to get our training tensors\ndef get_training_dataset(filenames, ordered = False):\n    dataset = load_dataset(filenames, ordered = ordered)\n    dataset = dataset.map(data_augment, num_parallel_calls = AUTO)\n    dataset = dataset.map(arcface_format, num_parallel_calls = AUTO)\n    dataset = dataset.map(lambda posting_id, image, label_group, matches: (image, label_group))\n    dataset = dataset.repeat()\n    dataset = dataset.shuffle(2048)\n    dataset = dataset.batch(config.BATCH_SIZE)\n    dataset = dataset.prefetch(AUTO)\n    return dataset\n\n# Function to count how many photos we have in\ndef count_data_items(filenames):\n    # Assuming 200000 Images per TFRecord\n    if config.dataset=='v2':\n        return 200000*len(filenames)\n    else:\n        n = [int(re.compile(r\"-([0-9]*)\\.\").search(filename).group(1)) for filename in filenames]\n        return np.sum(n)\n\nNUM_TRAINING_IMAGES = count_data_items(TRAINING_FILENAMES)\nprint(f'Dataset: {NUM_TRAINING_IMAGES} training images')","metadata":{"id":"RjpDnZoku-_D","outputId":"51bc6b2f-6a07-4629-fe1f-dba463b4009e","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"row = 10; col = 8;\nrow = min(row,config.BATCH_SIZE//col)\nN_TRAIN = count_data_items(TRAINING_FILENAMES)\nprint(N_TRAIN)\nds = get_training_dataset(TRAINING_FILENAMES, ordered = False)\n\nfor (sample,label) in ds:\n    img = sample['inp1']\n    plt.figure(figsize=(25,int(25*row/col)))\n    for j in range(row*col):\n        plt.subplot(row,col,j+1)\n        plt.title(label[j].numpy())\n        plt.axis('off')\n        plt.imshow(img[j,])\n    plt.show()\n    break\nprint(img.shape)","metadata":{"id":"9ZLvqUs4gdz6","outputId":"c8010a08-f10a-4ef6-8adb-64ac5a00abfc","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Model","metadata":{"id":"6xl8qKbRgdz6"}},{"cell_type":"code","source":"# Arcmarginproduct class keras layer\nclass 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":{"id":"spUWeq_2gdz7","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#CurricularFace class keras layer\nclass CurricularFace(tf.keras.layers.Layer):\n    '''\n    Implements Curricular Face.\n\n    Reference:\n        To be added\n    '''\n    def __init__(self, n_classes, s=30, m=0.50, easy_margin=False,\n                 ls_eps=0.0, **kwargs):\n\n        super(CurricularFace, 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(CurricularFace, self).build(input_shape[0])\n\n        self.kernel = 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            aggregation=tf.VariableAggregation.MEAN,\n            regularizer=None)\n        \n        self.t = self.add_weight(\n              name='t',\n              shape=(1,),\n              dtype='float32',\n              initializer=tf.keras.initializers.get('zeros'),\n              synchronization=tf.VariableSynchronization.ON_READ,\n              trainable=False,\n              aggregation=tf.VariableAggregation.MEAN,\n              experimental_autocast=False)\n    \n    def _assign_new_value(self, variable, value):\n        with tf.keras.backend.name_scope('AssignNewValue') as scope:\n            if tf.compat.v1.executing_eagerly_outside_functions():\n                return variable.assign(value, name=scope)\n            else:\n                with tf.compat.v1.colocate_with(variable):  # pylint: disable=protected-access\n                    return tf.compat.v1.assign(variable, value, name=scope)\n\n    def call(self, inputs):\n        embeddings, label = inputs\n        label = tf.cast(label,tf.int32)\n        embeddings = tf.math.l2_normalize(embeddings, axis=1)\n        kernel_norm = tf.math.l2_normalize(self.kernel, axis=0)\n        cos_theta = tf.matmul(embeddings, kernel_norm)\n        cos_theta = tf.clip_by_value(cos_theta,-1,1) # for numerical stability\n        origin_cos = tf.identity(cos_theta)\n        out = tf.stack([tf.range(0,tf.shape(embeddings)[0]),label],axis=1)\n        target_logit = tf.gather_nd(cos_theta,tf.transpose([tf.range(0,tf.shape(embeddings)[0]),label]))\n\n        sin_theta = tf.math.sqrt(1.0 - tf.math.pow(target_logit, 2))\n        cos_theta_m = target_logit * self.cos_m - sin_theta * self.sin_m #cos(target+margin) \n        final_target_logit = tf.where(target_logit > self.th, cos_theta_m, target_logit - self.mm)\n        \n        \n        one_hot = tf.cast(\n                    tf.one_hot(label, depth=self.n_classes),\n                    dtype=cos_theta.dtype\n                )\n\n\n        mask = cos_theta > cos_theta_m[:,None]\n        hard_example = tf.identity(cos_theta)\n        self._assign_new_value(self.t, tf.reduce_mean(target_logit) * 0.01 + (1 - 0.01) * self.t)\n        cos_hard = hard_example * (self.t + hard_example)\n        cos_theta = tf.where(mask, cos_hard, cos_theta)  \n        cos_theta = one_hot*final_target_logit[:,None]+((1.0 - one_hot) * cos_theta)\n        output = cos_theta * self.s\n        return output\n","metadata":{"id":"jHLOiF6zgdz7","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"EFNS = [efn.EfficientNetB0, efn.EfficientNetB1, efn.EfficientNetB2, efn.EfficientNetB3, \n        efn.EfficientNetB4, efn.EfficientNetB5, efn.EfficientNetB6, efn.EfficientNetB7]\n\ndef freeze_BN(model):\n    # Unfreeze layers while leaving BatchNorm layers frozen\n    for layer in model.layers:\n        if not isinstance(layer, tf.keras.layers.BatchNormalization):\n            layer.trainable = True\n        else:\n            layer.trainable = False\n\n# Function to create our EfficientNetB3 model\ndef get_model():\n\n    if config.head=='arcface':\n        head = ArcMarginProduct\n    elif config.head=='curricularface':\n        head = CurricularFace\n    else:\n        assert 1==2, \"INVALID HEAD\"\n    \n    with strategy.scope():\n\n        margin = head(\n            n_classes = config.N_CLASSES, \n            s = 30, \n            m = 0.6, \n            name='head/arc_margin', \n            dtype='float32'\n            )\n\n        inp = tf.keras.layers.Input(shape = [config.IMAGE_SIZE, config.IMAGE_SIZE, 3], name = 'inp1')\n        label = tf.keras.layers.Input(shape = (), name = 'inp2')\n        \n        if config.model_type == 'effnetv1':\n            x = EFNS[config.EFF_NET](weights = 'noisy-student', include_top = False)(inp)\n            embed = tf.keras.layers.GlobalAveragePooling2D()(x)\n        elif config.model_type == 'effnetv2':\n            FEATURE_VECTOR = f'{EFFNETV2_ROOT}/tfhub_models/efficientnetv2-{config.EFF_NETV2}/feature_vector'\n            embed = tfhub.KerasLayer(FEATURE_VECTOR, trainable=True)(inp)\n            \n        embed = tf.keras.layers.Dropout(0.2)(embed)\n        embed = tf.keras.layers.Dense(512)(embed)\n        x = margin([embed, label])\n        \n        output = tf.keras.layers.Softmax(dtype='float32')(x)\n\n        model = tf.keras.models.Model(inputs = [inp, label], outputs = [output])\n        embed_model = tf.keras.models.Model(inputs = inp, outputs = embed)  \n        \n        opt = tf.keras.optimizers.Adam(learning_rate = config.LR)\n        if config.FREEZE_BATCH_NORM:\n            freeze_BN(model)\n\n        model.compile(\n            optimizer = opt,\n            loss = [tf.keras.losses.SparseCategoricalCrossentropy()],\n            metrics = [tf.keras.metrics.SparseCategoricalAccuracy()]\n            ) \n        \n        return model,embed_model","metadata":{"id":"WDFiBlOfgdz7","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_lr_callback(plot=False):\n    lr_start   = 0.000001\n    lr_max     = 0.000005 * config.BATCH_SIZE  \n    lr_min     = 0.000001\n    lr_ramp_ep = 4\n    lr_sus_ep  = 0\n    lr_decay   = 0.9\n   \n    def lrfn(epoch):\n        if config.RESUME:\n            epoch = epoch + config.RESUME_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(config.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":{"id":"-mqooKyQu-_G","outputId":"b916e064-b322-46da-84bd-8347041f01e4","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def is_interactive():\n    return 'runtime'    in get_ipython().config.IPKernelApp.connection_file\nIS_INTERACTIVE = is_interactive()\nprint(IS_INTERACTIVE)","metadata":{"id":"R4M69Qivu-_K","outputId":"e52f79b4-a23e-4a53-9271-463aeb5dff1f","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Snapshot(tf.keras.callbacks.Callback):\n    \n    def __init__(self,snapshot_min_epoch,fold):\n        super(Snapshot, self).__init__()\n        self.snapshot_min_epoch = snapshot_min_epoch\n        self.fold = fold\n        \n        \n    def on_epoch_end(self, epoch, logs=None):\n        # logs is a dictionary\n#         print(f\"epoch: {epoch}, train_acc: {logs['acc']}, valid_acc: {logs['val_acc']}\")\n        if epoch >=self.snapshot_min_epoch: # your custom condition         \n            self.model.save(config.save_dir+f\"/EF{config.EFF_NET}_fold{self.fold}_epoch{epoch}.h5\")\n        self.model.save(config.save_dir+f\"/EF{config.EFF_NET}_fold{self.fold}_last.h5\")","metadata":{"id":"W44pDoCeu-_K","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Training","metadata":{"id":"kavqry6Bgdz8"}},{"cell_type":"code","source":"seed_everything(config.SEED)\nVERBOSE = 1\ntrain_dataset = get_training_dataset(TRAINING_FILENAMES, ordered = False)\nSTEPS_PER_EPOCH = count_data_items(TRAINING_FILENAMES) // config.BATCH_SIZE\nif config.dataset=='v2':\n  STEPS_PER_EPOCH = STEPS_PER_EPOCH//3\ntrain_logger = tf.keras.callbacks.CSVLogger(config.save_dir+'/training-log-fold-%i.h5.csv'%config.FOLD_TO_RUN)\n# BUILD MODEL\nK.clear_session()\nmodel,embed_model = get_model()\nsnap = Snapshot(snapshot_min_epoch=config.SNAPSHOT_THRESHOLD,fold=config.FOLD_TO_RUN)\nmodel.summary()","metadata":{"id":"LZEQX4qSu-_L","outputId":"85207e90-1728-4701-b77e-5d29750f4ce0","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if config.RESUME:   \n  model.load_model(config.model_path+f\"EF{config.EFF_NET}_fold{config.FOLD_TO_RUN}_last.h5\")","metadata":{"id":"PZF23riq4MOj","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print('#### Image Size %i with EfficientNet B%i and batch_size %i'%\n      (config.IMAGE_SIZE,config.EFF_NET,config.BATCH_SIZE))\n\nhistory = model.fit(train_dataset,\n                steps_per_epoch = STEPS_PER_EPOCH,\n                epochs = config.EPOCHS,\n                callbacks = [snap,get_lr_callback(),train_logger], \n                verbose = VERBOSE)\n# model.save_weights(config.save_dir+'/fold-%i.h5')","metadata":{"id":"-u9PNeT3gdz8","outputId":"beb03cbc-17fd-4062-9e5d-d4707ef24deb","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Save Optimiser ","metadata":{"id":"QLOutLydgdz8"}},{"cell_type":"code","source":"symbolic_weights = getattr(model.optimizer, 'weights')\nweight_values = K.batch_get_value(symbolic_weights)\nwith open(config.save_dir+'/optimizer.pkl', 'wb') as f:\n    pickle.dump(weight_values, f)","metadata":{"id":"3cB1fMoUu-_M","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}