{"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 tensorflow as tf\ntry:\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    strategy = tf.distribute.get_strategy()\n\nAUTO = tf.data.experimental.AUTOTUNE\nprint(\"REPLICAS: \", strategy.num_replicas_in_sync)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from kaggle_datasets import KaggleDatasets","metadata":{},"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","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"save_dir = '.'\nEXPERIMENT = 0\nrun_ts = datetime.now().strftime('%Y%m%d-%H%M%S')\nprint(run_ts)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class config:\n     \n    SEED = 56\n    FOLD_TO_RUN = 20\n    FOLDS = 30\n    DEBUG = False\n    EVALUATE = True\n    RESUME = False\n    RESUME_EPOCH = None\n       \n    BATCH_SIZE = 8 * strategy.num_replicas_in_sync\n    IMAGE_SIZE = 768\n    N_CLASSES = 15587\n    model_type = 'effnetv1'  \n    EFF_NET = 7\n    EFF_NETV2 = 's-21k-ft1k'\n    FREEZE_BATCH_NORM = False\n    head = 'arcface' \n    EPOCHS = 31\n    LR = 0.001\n    message='baseline'\n    CUTOUT = False\n    save_dir = save_dir\n    KNN = 1000\n    \ndef count_data_items(filenames):\n    n = [int(re.compile(r\"-([0-9]*)\\.\").search(filename).group(1)) \n         for filename in filenames]\n    return np.sum(n)\n\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\nseed_everything(56)    ","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"MODEL_NAME = None\nif config.model_type == 'effnetv1':\n    MODEL_NAME = f'effnetv1_b{config.EFF_NET}'\nelif config.model_type == 'effnetv2':\n    MODEL_NAME = f'effnetv2_{config.EFF_NETV2}'\n\nconfig.MODEL_NAME = MODEL_NAME\nprint(MODEL_NAME)","metadata":{},"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":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"GCS_PATH1 = KaggleDatasets().get_gcs_path('happywhale-tfrecords-fullbody-768')\n    \ntrain_files = np.sort(np.array(tf.io.gfile.glob(GCS_PATH1 + '/train*.tfrec')))\ntest_files = np.sort(np.array(tf.io.gfile.glob(GCS_PATH1 + '/test*.tfrec')))\nprint(GCS_PATH1)\nprint(len(train_files),len(test_files),count_data_items(train_files),count_data_items(test_files))","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def arcface_format(posting_id, image, label_group, matches):\n    return posting_id, {'inp1': image, 'inp2': label_group}, label_group, matches\n\ndef arcface_inference_format(posting_id, image, label_group, matches):\n    return image,posting_id\n\ndef arcface_eval_format(posting_id, image, label_group, matches):\n    return image,label_group\n\ndef data_augment(posting_id, image, label_group, matches):\n    image = tf.image.random_flip_left_right(image, 6)\n    image = tf.image.random_hue(image, 0.01, 3)\n    image = tf.image.random_saturation(image, 0.65, 1.15)\n    image = tf.image.random_contrast(image, 0.7, 1.10, 4)\n    image = tf.image.random_brightness(image, 0.1, 5)\n    return posting_id, image, label_group, matches\n    \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\ndef read_labeled_tfrecord(example):\n    LABELED_TFREC_FORMAT = {'image_id': tf.io.FixedLenFeature([], tf.string),\n                            'image': tf.io.FixedLenFeature([], tf.string),\n                            'species': tf.io.FixedLenFeature([], tf.int64),\n                            'individual_id': tf.io.FixedLenFeature([], tf.int64),\n                            }\n    example = tf.io.parse_single_example(example, LABELED_TFREC_FORMAT)\n    posting_id = example['image_id']\n    image = decode_image(example['image'])\n    label_group = tf.cast(example['individual_id'], tf.int32)\n    matches = 1\n    return posting_id, image, label_group, matches\n\ndef load_dataset(filenames, ordered = False):\n    ignore_order = tf.data.Options()\n    if not ordered:\n        ignore_order.experimental_deterministic = False \n    dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads = AUTO)\n    dataset = dataset.with_options(ignore_order)\n    dataset = dataset.map(read_labeled_tfrecord, num_parallel_calls = AUTO) \n    return dataset\n\ndef get_training_dataset(filenames):\n    dataset = load_dataset(filenames, ordered = False)\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\ndef get_val_dataset(filenames):\n    dataset = load_dataset(filenames, ordered = True)\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.batch(config.BATCH_SIZE)\n    dataset = dataset.prefetch(AUTO)\n    return dataset\n\ndef get_eval_dataset(filenames, get_targets = True):\n    dataset = load_dataset(filenames, ordered = True)\n    dataset = dataset.map(data_augment, num_parallel_calls = AUTO)\n    dataset = dataset.map(arcface_eval_format, num_parallel_calls = AUTO)\n    if not get_targets:\n        dataset = dataset.map(lambda image, target: image)\n    dataset = dataset.batch(config.BATCH_SIZE)\n    dataset = dataset.prefetch(AUTO)\n    return dataset\n\ndef get_test_dataset(filenames, get_names = True):\n    dataset = load_dataset(filenames, ordered = True)\n    dataset = dataset.map(data_augment, num_parallel_calls = AUTO)\n    dataset = dataset.map(arcface_inference_format, num_parallel_calls = AUTO)\n    if not get_names:\n        dataset = dataset.map(lambda image, posting_id: image)\n    dataset = dataset.batch(config.BATCH_SIZE)\n    dataset = dataset.prefetch(AUTO)\n    return dataset","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class ArcMarginProduct(tf.keras.layers.Layer):\n    \n    def __init__(self, n_classes, s=30, m=0.50, easy_margin=False, ls_eps=0.0, **kwargs):\n        super(ArcMarginProduct, self).__init__(**kwargs)\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_classes': self.n_classes,'s': self.s, 'm': self.m,\n                       'ls_eps': self.ls_eps, 'easy_margin': self.easy_margin, })\n        return config\n\n    def build(self, input_shape):\n        super(ArcMarginProduct, self).build(input_shape[0])\n        self.W = self.add_weight(name='W', shape=(int(input_shape[0][-1]), self.n_classes),\n                                 initializer='glorot_uniform', dtype='float32',\n                                 trainable=True, regularizer=None)\n    def call(self, inputs):\n        X, y = inputs\n        y = tf.cast(y, dtype=tf.int32)\n        cosine = tf.matmul(tf.math.l2_normalize(X, axis=1), tf.math.l2_normalize(self.W, axis=0))\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(tf.one_hot(y, depth=self.n_classes), dtype=cosine.dtype)\n        if self.ls_eps > 0:\n            one_hot = (1 - self.ls_eps) * one_hot + self.ls_eps / self.n_classes\n        output = (one_hot * phi) + ((1.0 - one_hot) * cosine)\n        output *= self.s\n        return output\n\nEFNS = [efn.EfficientNetB0, efn.EfficientNetB1, efn.EfficientNetB2, efn.EfficientNetB3, \n        efn.EfficientNetB4, efn.EfficientNetB5, efn.EfficientNetB6, efn.EfficientNetB7]\n\ndef freeze_BN(model):\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\ndef get_model():\n    if config.head=='arcface':\n        head = ArcMarginProduct\n    else:\n        assert 1==2, \"INVALID HEAD\"\n    \n    with strategy.scope():\n        margin = head(n_classes = config.N_CLASSES, s=30, m=0.5, name=f'head/{config.head}', dtype='float32')\n\n        label = tf.keras.layers.Input(shape = (), name = 'inp2')\n        \n        if config.model_type == 'effnetv1':\n            inp = EFNS[config.EFF_NET](weights='noisy-student', include_top=False,\n                                       input_shape = [config.IMAGE_SIZE, config.IMAGE_SIZE, 3])\n            inp.layers[0]._name = 'inp1'\n            x1=tf.keras.layers.GlobalAveragePooling2D()(inp.layers[-1].output)\n            x2=tf.keras.layers.GlobalAveragePooling2D()(inp.layers[-5].output)\n            x3=tf.keras.layers.GlobalAveragePooling2D()(inp.layers[-7].output)\n            x4=tf.keras.layers.GlobalAveragePooling2D()(inp.layers[-13].output)\n            embed =  tf.concat([x1,x2,x3,x4],axis = 1)    \n            \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.3)(embed)\n        embed = tf.keras.layers.Dense(2048)(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.input, label], outputs = [output])\n        embed_model = tf.keras.models.Model(inputs = inp.input, 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(optimizer = opt, loss = [tf.keras.losses.SparseCategoricalCrossentropy()],\n                      metrics = [tf.keras.metrics.SparseCategoricalAccuracy(),\n                                 tf.keras.metrics.SparseTopKCategoricalAccuracy(k=5)]) \n        return model, embed_model\n","metadata":{},"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        else:\n            lr = (lr_max - lr_min) * lr_decay**(epoch - lr_ramp_ep - lr_sus_ep) + lr_min\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":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Snapshot(tf.keras.callbacks.Callback):\n    \n    def __init__(self,fold,snapshot_epochs=[]):\n        super(Snapshot, self).__init__()\n        self.snapshot_epochs = snapshot_epochs\n        self.fold = fold\n          \n    def on_epoch_end(self, epoch, logs=None):\n        if epoch in self.snapshot_epochs:     \n            self.model.save_weights(config.save_dir+f\"/EF{config.MODEL_NAME}_epoch{epoch}.h5\")\n        self.model.save_weights(config.save_dir+f\"/{config.MODEL_NAME}_last.h5\")\n\nTRAINING_FILENAMES = [x for i,x in enumerate(train_files) if i%config.FOLDS!=config.FOLD_TO_RUN]\nVALIDATION_FILENAMES = [x for i,x in enumerate(train_files) if i%config.FOLDS==config.FOLD_TO_RUN]\nprint(len(TRAINING_FILENAMES),len(VALIDATION_FILENAMES),count_data_items(TRAINING_FILENAMES),\n      count_data_items(VALIDATION_FILENAMES))","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#seed_everything(config.SEED)\nVERBOSE = 1\ntrain_dataset = get_training_dataset(TRAINING_FILENAMES)\nval_dataset = get_val_dataset(VALIDATION_FILENAMES)\nSTEPS_PER_EPOCH = count_data_items(TRAINING_FILENAMES) // config.BATCH_SIZE\ntrain_logger = tf.keras.callbacks.CSVLogger(config.save_dir+'/training-log-fold-%i.h5.csv'%config.FOLD_TO_RUN)\nsv_loss = tf.keras.callbacks.ModelCheckpoint(\n    config.save_dir+f\"/{config.MODEL_NAME}_loss.h5\", monitor='val_loss', verbose=0, save_best_only=True,\n    save_weights_only=True, mode='min', save_freq='epoch')\nK.clear_session()\nmodel,embed_model = get_model()\nsnap = Snapshot(fold=config.FOLD_TO_RUN,snapshot_epochs=[13,27])\n#model.summary()\n\nif config.RESUME:   \n    model.load_weights(config.resume_model_wts)","metadata":{},"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, validation_data = val_dataset, steps_per_epoch = STEPS_PER_EPOCH,\n                    epochs = config.EPOCHS, callbacks = [snap,get_lr_callback(),train_logger,sv_loss], \n                    verbose = VERBOSE)","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}