{"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":"# CurricularFace","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19"}},{"cell_type":"code","source":"class CurricularFace(tf.keras.layers.Layer):\n    def __init__(self, n_classes, s=30, m=0.50, easy_margin=False,\n                 ls_eps=0.0, **kwargs):\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        self._USE_V2_BEHAVIOR = True\n\n    def _assign_new_value(self, variable, value):\n        with 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\n    def _get_training_value(self, training=None):\n        if training is None:\n          training = backend.learning_phase()\n        if self._USE_V2_BEHAVIOR:\n          if isinstance(training, int):\n            training = bool(training)\n          if not self.trainable:\n            # When the layer is not trainable, it overrides the value passed from\n            # model.\n            training = False\n        return training\n\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.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        self.t = self.add_weight(\n            name='t',\n            shape=(1),\n            initializer=tf.zeros_initializer(),\n            dtype='float32',\n            trainable=False,\n            regularizer=None,\n            aggregation=tf.VariableAggregation.MEAN,\n            experimental_autocast=False,\n            synchronization=tf.VariableSynchronization.ON_READ)\n        \n    def call(self, inputs, training=None):\n        X, y = inputs\n        y = tf.cast(y, dtype=tf.int32)\n\n        do_training = self._get_training_value(training)\n\n        if do_training:\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\n            target_logit = tf.reduce_sum(cosine * tf.cast(tf.one_hot(y, depth=self.n_classes),dtype=cosine.dtype), axis=-1)\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\n\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        \n            t = tf.reduce_mean(target_logit) * 0.01 + (1 - 0.01) * self.t\n            self._assign_new_value(self.t, t)\n            cosine = tf.where(cosine > tf.expand_dims(cos_theta_m, axis=-1), cosine*(self.t+cosine), cosine)\n\n            if self.ls_eps > 0:\n                one_hot = (1 - self.ls_eps) * one_hot + self.ls_eps / self.n_classes\n\n            output = (one_hot * phi) + ((1.0 - one_hot) * cosine)\n            output *= self.s\n\n        else:\n            output = tf.matmul(\n                tf.math.l2_normalize(X, axis=1),\n                tf.math.l2_normalize(self.W, axis=0)\n            )\n\n        return output","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Gradient Accumulation","metadata":{}},{"cell_type":"code","source":"EPOCHS = 8\nN_TRAIN = 1_500_000\nSTEPS_PER_TPU_CALL = 1\nAVG_N_BATCHES = 16\nBATCH_SIZE = 8*REPLICAS\nSTEPS_PER_EPOCH = N_TRAIN//STEPS_PER_TPU_CALL//AVG_N_BATCHES//BATCH_SIZE\n\n###########################################\n###### Create objects and Initialize ######\n###########################################\n\nwith strategy.scope():\n    model = build_model()\n    loss_fn = tf.keras.losses.sparse_categorical_crossentropy\n    optimizer=tf.keras.optimizers.Adam(learning_rate=0.0001)\n    train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()\n\n###########################################\n##### Function to calculate gradients #####\n###########################################\n\n@tf.function\ndef get_gradients(I,O):\n    with tf.GradientTape() as tape:\n        probabilities = model(I, training=True)\n        loss = loss_fn(O, probabilities)\n    grads = tape.gradient(loss, model.trainable_variables)\n    train_accuracy.update_state(O, probabilities)\n    train_acc = train_accuracy.result()\n    return grads\n\n###########################################\n####### Function to apply gradients #######\n###########################################\n\n@tf.function\ndef apply_gradients(grads):\n    optimizer.apply_gradients(zip(grads, model.trainable_variables))\n\n\ndef _minimum_control_deps(outputs):\n  \"\"\"Returns the minimum control dependencies to ensure step succeeded.\"\"\"\n  if tf.executing_eagerly():\n    return []  # Control dependencies not needed.\n  outputs = tf.nest.flatten(outputs, expand_composites=True)\n  for out in outputs:\n    # Variables can't be control dependencies.\n    if not isinstance(out, tf.Variable):\n      return [out]  # Return first Tensor or Op from outputs.\n  return []\n\n############################################\n########### Train Step Function ############\n############################################\n\n@tf.function\ndef train_step(train_data_iter):\n    grads = strategy.run(get_gradients, next(train_data_iter))\n    train_acc = train_accuracy.result()\n\n    with tf.device('/TPU:0'):\n        grads = strategy.reduce('mean', grads, axis=None)\n        grads = [g/(AVG_N_BATCHES*1.0) for g in grads]\n\n    grads_0 = strategy.run(get_gradients, next(train_data_iter))\n    train_acc = train_accuracy.result()\n\n    with tf.device('/TPU:0'):\n        grads_0 = strategy.reduce('mean', grads_0, axis=None)\n        grads = [g0+g1/(AVG_N_BATCHES*1.0) for g0,g1 in zip(grads,grads_0)]\n\n    for _ in tf.range(AVG_N_BATCHES-2):\n\n        with tf.control_dependencies(_minimum_control_deps(grads)):\n            grads_0 = strategy.run(get_gradients, next(train_data_iter))\n            with tf.device('/TPU:0'):\n                grads_0 = strategy.reduce('mean', grads_0, axis=None)\n                grads = [g0+g1/(AVG_N_BATCHES*1.0) for g0,g1 in zip(grads,grads_0)]\n            \n        train_acc = train_accuracy.result()\n    with tf.control_dependencies(_minimum_control_deps(grads)):\n        strategy.run(apply_gradients, args = (grads,))\n\n# distributed dataset\ntrain_dist_ds = strategy.experimental_distribute_dataset(get_dataset(files_train,batch_size = BATCH_SIZE,mode='train'))\n\n# dataset iterator\ntrain_data_iter = iter(train_dist_ds)\n\n# custom training loop\nstart = time.time()\nfor epoch in range(EPOCHS):\n    pbar = tqdm(range(STEPS_PER_EPOCH+1))\n    for steps in pbar:\n        train_step(train_data_iter)\n        train_acc = train_accuracy.result().numpy()\n        pbar.set_description('Train Accuracy: '+str(train_acc))\n    \n    print(f'\\n\\nEVALUATING EPOCH: {epoch}')\n    model.evaluate(get_dataset(files_valid,mode='val'))\n    print(' \\n\\n')\n    train_accuracy.reset_states()\n    model.save_weights(BASE_SAVE_DIR+f'/l-512-custom-epoch-{epoch:02d}.h5')","metadata":{},"execution_count":null,"outputs":[]}]}