{"cells":[{"metadata":{},"cell_type":"markdown","source":"### In this notebook, I will show you Tensorflow implementation of ArcFace layer, GeM pooling layer and how to train everything on TPU.\n\n#### TPU must read and store data from/to Google cloud storage. Although you could access input data in this gs dir: `KaggleDatasets().get_gcs_path()`, I would recommand you to shard and serialise the data beforehand using [this script](https://github.com/tensorflow/models/blob/master/research/delf/delf/python/training/build_image_dataset.py) from DELF repo. Afterwards you should upload shards to your google cloud storage and uses `tf.data.TFRecordDataset` to load them. It can increase your training time by 4x. Here I will continue as data has been sharded and serialised.  \n","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"### Hyper params","execution_count":null},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import math\n\nEPOCHS = 10\nSTEPS_PER_TPU_CALL = 1  # set this number higher for less TPU idle time\nbatch_size = 16\nimage_size = 500\nlearning_rate = 1e-5  # should be smaller than training on single GPU\nfeature_size = 2048  # Embedding size before the output layer\n\n# ArcFace params\nmargin = 0.1  # DELG used 0.1, original ArcFace paper used 0.5. When margin is 0, it should be the same as doing a normal softmax but with embedding and weight normalised.\nlogit_scale = int(math.sqrt(feature_size))\n\n# GeM params\ngem_p = 3.\ntrain_p = False  # whether to learn gem_p or not\n","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"# Google storage settings\n# Set your own project id here\nPROJECT_ID = 'YOUR_PROJECT_ID'\nfrom google.cloud import storage\nstorage_client = storage.Client(project=PROJECT_ID)\n\n# Step 1: Get the credential from the Cloud SDK\n# Get the credential from the Cloud SDK\nfrom kaggle_secrets import UserSecretsClient\nuser_secrets = UserSecretsClient()\nuser_credential = user_secrets.get_gcloud_credential()\n\n# Step 2: Set the credentials\n# Set the credentials\nuser_secrets.set_tensorflow_credential(user_credential)\n\n# Step 3: Take note of the GCS path\n# Take note of the GCS path\nfrom kaggle_datasets import KaggleDatasets\ndata_dir = KaggleDatasets().get_gcs_path()\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Init TPU","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import pandas as pd\nimport functools\nfrom tqdm import tqdm\n\nimport tensorflow as tf\nfrom pathlib import Path\nimport os\n\n# detect and init the TPU\ntpu = tf.distribute.cluster_resolver.TPUClusterResolver()\ntf.config.experimental_connect_to_cluster(tpu)\ntf.tpu.experimental.initialize_tpu_system(tpu)\n\n# instantiate a distribution strategy\ntpu_strategy = tf.distribute.experimental.TPUStrategy(tpu)\n\n# instantiating the model in the strategy scope creates the model on the TPU\n\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Dataset","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"def _parse_example(example, name_to_features, image_size, augmentation, unique_landmark_ids):\n    parsed_example = tf.io.parse_single_example(example, name_to_features)\n    # Parse to get image.\n    image = parsed_example['image/encoded']\n    image = tf.io.decode_jpeg(image)\n    image = tf.cast(image, tf.float32)\n    image = tf.math.divide(tf.subtract(image, 128.0), 128.0)\n    # if augmentation:\n    # image = _ImageNetCrop(image)\n    # else:\n    image = tf.image.resize(image, [image_size, image_size])\n    image.set_shape([image_size, image_size, 3])\n    # Parse to get label.\n    label = parsed_example['image/class/label']\n    label = tf.reduce_min(tf.where(tf.equal(unique_landmark_ids, label)))\n\n    return image, label\n\n\ndef create_dataset(file_pattern, unique_landmark_ids, augmentation: bool = False):\n    AUTO = tf.data.experimental.AUTOTUNE\n    ignore_order = tf.data.Options()\n    ignore_order.experimental_deterministic = False\n    filenames = tf.io.gfile.glob(file_pattern)\n\n    dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=AUTO).shuffle(1000)\n\n    # Create a description of the features.\n    feature_description = {\n        'image/height': tf.io.FixedLenFeature([], tf.int64, default_value=0),\n        'image/width': tf.io.FixedLenFeature([], tf.int64, default_value=0),\n        'image/channels': tf.io.FixedLenFeature([], tf.int64, default_value=0),\n        'image/format': tf.io.FixedLenFeature([], tf.string, default_value=''),\n        'image/id': tf.io.FixedLenFeature([], tf.string, default_value=''),\n        'image/filename': tf.io.FixedLenFeature([], tf.string, default_value=''),\n        'image/encoded': tf.io.FixedLenFeature([], tf.string, default_value=''),\n        'image/class/label': tf.io.FixedLenFeature([], tf.int64, default_value=0),\n    }\n\n    customized_parse_func = functools.partial(\n        _parse_example,\n        name_to_features=feature_description,\n        image_size=image_size,\n        augmentation=augmentation,\n        unique_landmark_ids=unique_landmark_ids,\n    )\n    dataset = dataset.with_options(ignore_order)\n    dataset = dataset.map(customized_parse_func)\n    dataset = dataset.batch(batch_size)\n    return dataset","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Define Model and layer","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import tensorflow as tf\nimport functools\nimport math\n\n\nclass GeMPoolingLayer(tf.keras.layers.Layer):\n    def __init__(self, p=1., train_p=False):\n        super().__init__()\n        if train_p:\n            self.p = tf.Variable(p, dtype=tf.float32)\n        else:\n            self.p = p\n        self.eps = 1e-6\n\n    def call(self, inputs: tf.Tensor, **kwargs):\n        inputs = tf.clip_by_value(inputs, clip_value_min=1e-6, clip_value_max=tf.reduce_max(inputs))\n        inputs = tf.pow(inputs, self.p)\n        inputs = tf.reduce_mean(inputs, axis=[1, 2], keepdims=False)\n        inputs = tf.pow(inputs, 1./self.p)\n        return inputs\n\n\nclass DelfArcFaceModel(tf.keras.Model):\n    def __init__(self, input_shape, n_classes, margin, logit_scale, feature_size, p=None, train_p=False):\n        super().__init__()\n        self.backbone = tf.keras.applications.ResNet50(include_top=False, weights=\"imagenet\", input_shape=input_shape)\n        self.backbone.summary()\n\n        if p is not None:\n            self.global_pooling = GeMPoolingLayer(p, train_p=train_p)\n        else:\n            self.global_pooling = functools.partial(tf.reduce_mean, axis=[1, 2], keepdims=False)\n        self.dense1 = tf.keras.layers.Dense(feature_size, activation=None, kernel_initializer=\"glorot_normal\")\n        # self.bn1 = tf.keras.layers.BatchNormalization()\n        self.arcface = ArcFaceLayer(n_classes, margin, logit_scale)\n\n    def call(self, inputs, training=True, mask=None):\n        images, labels = inputs\n        x = self.extract_feature(images)\n        x = self.arcface((x, labels))\n        return x\n\n    def extract_feature(self, inputs):\n        x = self.backbone(inputs)\n        x = self.global_pooling(x)\n        x = self.dense1(x)\n        return x\n\n\nclass ArcFaceLayer(tf.keras.layers.Layer):\n    def __init__(self, num_classes, margin, logit_scale):\n        super().__init__()\n        self.num_classes = num_classes\n        self.margin = margin\n        self.logit_scale = logit_scale\n\n    def build(self, input_shape):\n        self.w = self.add_weight(\"weights\", shape=[int(input_shape[0][-1]), self.num_classes], initializer=tf.keras.initializers.get(\"glorot_normal\"))\n        self.cos_m = tf.identity(tf.cos(self.margin), name='cos_m')\n        self.sin_m = tf.identity(tf.sin(self.margin), name='sin_m')\n        self.th = tf.identity(tf.cos(math.pi - self.margin), name='th')\n        self.mm = tf.multiply(self.sin_m, self.margin, name='mm')\n\n    def call(self, inputs, training=True, mask=None):\n        embeddings, labels = inputs\n        normed_embeddings = tf.nn.l2_normalize(embeddings, 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_embeddings, normed_w, name='cos_t')\n        sin_t = tf.sqrt(1. - cos_t ** 2, name='sin_t')\n\n        cos_mt = tf.subtract(cos_t * self.cos_m, sin_t * self.sin_m, name='cos_mt')\n\n        cos_mt = tf.where(cos_t > self.th, cos_mt, cos_t - self.mm)\n\n        mask = tf.one_hot(tf.cast(labels, tf.int32), depth=self.num_classes,\n                          name='one_hot_mask')\n\n        logits = tf.where(mask == 1., cos_mt, cos_t)\n        logits = tf.multiply(logits, self.logit_scale, 'arcface_logist')\n        return logits\n\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Distributed tf training and test function","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"\ndef train_step(images, labels):\n    with tf.GradientTape() as tape:\n        # training=True is only needed if there are layers with different\n        # behavior during training versus inference (e.g. Dropout).\n        predictions = model((images, labels), training=True)\n\n        loss = compute_loss(labels, predictions)\n    gradients = tape.gradient(loss, model.trainable_variables)\n    optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n\n    train_loss(loss * tpu_strategy.num_replicas_in_sync)\n    train_accuracy(labels, predictions)\n    return loss\n\n\n@tf.function\ndef distributed_train_steps(training_set_iter, steps_per_call):\n    for _ in tf.range(steps_per_call):\n        per_replica_losses = tpu_strategy.run(train_step, next(training_set_iter))\n    # return tpu_strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)\n\n\ndef test_step(images, labels):\n    # training=False is only needed if there are layers with different\n    # behavior during training versus inference (e.g. Dropout).\n    predictions = model(images, training=False)\n    t_loss = loss_object(labels, predictions)\n\n    test_loss(t_loss)\n    test_accuracy(labels, predictions)\n\n\n@tf.function\ndef distributed_test_step(images, labels):\n    return tpu_strategy.run(test_step, args=(images, labels, ))\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"\n### Main Training script\n","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"checkpoint_dir = \"gs://{YOUR_GS_BUCKET}/checkpoints\"\ntrain_tf_records_dir = \"gs://{YOUR_GS_BUCKET}/{YOUR_DATA_SHARD_DIR}/train*\"\ntest_tf_records_dir = \"gs://{YOUR_GS_BUCKET}/{YOUR_DATA_SHARD_DIR}/validation*\"\n\n\ntraining_csv_path = os.path.join(data_dir, \"train.csv\")\ntrain_csv = pd.read_csv(str(training_csv_path))\nnum_samples = len(train_csv[\"id\"].tolist())\nunique_landmark_ids = train_csv[\"landmark_id\"].unique().tolist()\nunique_landmark_ids = tf.convert_to_tensor(unique_landmark_ids, dtype=tf.int64)\n\ntraining_set = create_dataset(train_tf_records_dir, unique_landmark_ids)\ntraining_set = tpu_strategy.experimental_distribute_dataset(training_set)\n\ntest_set = create_dataset(test_tf_records_dir, unique_landmark_ids)\ntest_set = tpu_strategy.experimental_distribute_dataset(test_set)\n\nwith tpu_strategy.scope():\n    model = DelfArcFaceModel(\n        input_shape=(image_size, image_size, 3), n_classes=len(unique_landmark_ids), margin=margin, logit_scale=logit_scale,\n        p=gem_p, train_p=train_p, feature_size=feature_size\n    )\n    optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)\n    loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=tf.keras.losses.Reduction.NONE)\n    train_loss = tf.keras.metrics.Mean(name='train_loss')\n    train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')\n    test_loss = tf.keras.metrics.Mean(name='test_loss')\n    test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')\n\n    def compute_loss(labels, predictions):\n        per_example_loss = loss_object(labels, predictions)\n        return tf.nn.compute_average_loss(per_example_loss, global_batch_size=batch_size)\n\ntraining_set_iter = iter(training_set)\n\nfor epoch in range(EPOCHS):\n    # Reset the metrics at the start of the next epoch\n    train_loss.reset_states()\n    train_accuracy.reset_states()\n    test_loss.reset_states()\n    test_accuracy.reset_states()\n    step = 0\n\n    with tqdm(total=int(num_samples)*0.8) as pbar:\n        while True:\n            distributed_train_steps(training_set_iter, tf.convert_to_tensor(STEPS_PER_TPU_CALL))\n            template = 'Epoch {}, Training, Loss: {:.4f}, Accuracy: {:.4f}'\n            pbar.set_description(template.format(epoch + 1, train_loss.result(), train_accuracy.result() * 100))\n            if step % save_interval == 0:\n                if step == 0:\n                    model.summary()\n                    print()\n                    print(\"\\nlearning rate: {}\\nmargin: {}\\nlogit_scale: {}\\ngem_p: {}\\ntrain_p{}\\n\".format(learning_rate, margin, logit_scale, gem_p, train_p))\n\n                checkpoint_path = str(os.path.join(checkpoint_dir, \"cp_epoch_{}_step_{}\".format(epoch, step)))\n                model.save_weights(checkpoint_path)\n                print(\"Model saved to {}\".format(checkpoint_path))\n            step += batch_size * STEPS_PER_TPU_CALL\n            pbar.update(batch_size * STEPS_PER_TPU_CALL)\n            if step >= int(num_samples)*0.8:\n                break\n\n    with tqdm(total=int(num_samples)*0.2) as pbar:\n        for test_images, test_labels in test_set:\n            distributed_test_step(test_images, test_labels)\n            template = 'Epoch {}, Validation, Loss: {:.4f}, Accuracy: {:.4f}'\n            pbar.set_description(template.format(epoch + 1, test_loss.result(), test_accuracy.result() * 100))\n            pbar.update(batch_size)\n\n    template = 'Epoch {}, \\nTraining Loss: {}, Accuracy: {}\\nTest Loss: {}, Accuracy: {}'\n    print(template.format(epoch + 1, train_loss.result(), train_accuracy.result() * 100, test_loss.result(), test_accuracy.result() * 100))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"\nSource:\n\nArcFace: \n\nhttps://github.com/auroua/InsightFace_TF/blob/master/losses/face_losses.py\n\nGeM: \n\nhttps://github.com/filipradenovic/cnnimageretrieval-pytorch/blob/master/cirtorch/layers/functional.py\n\nTPU: \n\nhttps://www.tensorflow.org/guide/tpu#train_a_model_using_custom_training_loop\n\nhttps://www.tensorflow.org/tutorials/distribute/custom_training#training_loop","execution_count":null}],"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":4,"nbformat_minor":4}