{"cells":[{"metadata":{},"cell_type":"markdown","source":"## Train Knowledge Distillation model\nThis notebook's idea is alomost same as 3 stage training's step2 in [RANZCR / ResNet200D / 3-stage training / step2](https://www.kaggle.com/yasufuminakama/ranzcr-resnet200d-3-stage-training-step2).\n\nKeras version implementation is borrowed from Keras official \"Knowledge Distillation\" example code.\n[Knowledge Distillation]\n(https://keras.io/examples/vision/knowledge_distillation/#construct-distiller-class)\n\n   \n\nTeacher model training notebook:  \n[[Keras TPU] RANZCR Train annotation](https://www.kaggle.com/enukuro/keras-tpu-ranzcr-train-annotation)\n\n\nCreatting annotation tfrecords notebook:   \n[Annotation RANZCR CLiP 900](https://www.kaggle.com/enukuro/annotation-ranzcr-clip-900)"},{"metadata":{"id":"TB205EIUGzHu","outputId":"d88358d9-08cc-42c8-dfc2-4edf7fa1eeed","trusted":true},"cell_type":"code","source":"import os\n\nimport numpy as np\nimport pandas as pd\nimport random\nimport tensorflow as tf\nimport tensorflow.keras.layers as L\nimport tensorflow_addons as tfa\nfrom keras import backend as K\nfrom kaggle_datasets import KaggleDatasets\nfrom keras.applications.xception import Xception as BaseModel\nfrom keras.applications.xception import preprocess_input\nimport itertools\nimport gc","execution_count":null,"outputs":[]},{"metadata":{"id":"nUlL507YG1G0","outputId":"4ab30bb3-ec02-495f-817a-44a68ec4602d","trusted":true},"cell_type":"code","source":"import tensorflow as tf\nprint(\"Tensorflow version \" + tf.__version__)\n\ntry:\n    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\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.experimental.TPUStrategy(tpu)\nelse:\n    strategy = tf.distribute.get_strategy()\n\nAUTOTUNE = tf.data.experimental.AUTOTUNE","execution_count":null,"outputs":[]},{"metadata":{"id":"1kOVhjvwG23G","trusted":true},"cell_type":"code","source":"SEED = 42\n\ndef seed_everything(seed):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    tf.random.set_seed(seed)\n    \nseed_everything(SEED)\n\nBATCH_SIZE = strategy.num_replicas_in_sync * 16\nIMG_SIZE = 900\nNUM_CLASSES = 11\n\ntarget_fold = 1","execution_count":null,"outputs":[]},{"metadata":{"id":"XJ3-_YipG53o","trusted":true},"cell_type":"code","source":"feature_map = {\n    'image': tf.io.FixedLenFeature([], tf.string),\n    'image_annotation': tf.io.FixedLenFeature([], tf.string),\n    'StudyInstanceUID': tf.io.FixedLenFeature([], tf.string),  \n    'ETT - Abnormal': tf.io.FixedLenFeature([], tf.int64),\n    'ETT - Borderline': tf.io.FixedLenFeature([], tf.int64),\n    'ETT - Normal': tf.io.FixedLenFeature([], tf.int64),\n    'NGT - Abnormal': tf.io.FixedLenFeature([], tf.int64),\n    'NGT - Borderline': tf.io.FixedLenFeature([], tf.int64),\n    'NGT - Incompletely Imaged': tf.io.FixedLenFeature([], tf.int64),\n    'NGT - Normal': tf.io.FixedLenFeature([], tf.int64),\n    'CVC - Abnormal': tf.io.FixedLenFeature([], tf.int64),\n    'CVC - Borderline': tf.io.FixedLenFeature([], tf.int64),\n    'CVC - Normal': tf.io.FixedLenFeature([], tf.int64),\n    'Swan Ganz Catheter Present': tf.io.FixedLenFeature([], tf.int64)}\n\ndef decode_image(image_data):\n    image = tf.image.decode_jpeg(image_data, channels=3)\n    image = tf.image.resize(image, [IMG_SIZE, IMG_SIZE])\n    image = tf.reshape(image, [IMG_SIZE, IMG_SIZE, 3])\n    return tf.cast(image, tf.float32)\n\ndef read_tfrecord(example):\n    example = tf.io.parse_single_example(example, feature_map)\n    image = decode_image(example['image'])\n    image_annotation = decode_image(example['image_annotation'])\n    target = [\n        example['ETT - Abnormal'],\n        example['ETT - Borderline'],\n        example['ETT - Normal'],\n        example['NGT - Abnormal'],\n        example['NGT - Borderline'],\n        example['NGT - Incompletely Imaged'],\n        example['NGT - Normal'],\n        example['CVC - Abnormal'],\n        example['CVC - Borderline'],\n        example['CVC - Normal'],\n        example['Swan Ganz Catheter Present']]\n    return [preprocess_input(image), preprocess_input(image_annotation)], tf.cast(target, tf.float32)\n\ndef data_augment(img, target):\n    img = tf.map_fn(lambda x: tf.image.random_flip_left_right(x), img)\n    return img, target\n\ndef get_dataset(filenames, shuffled=False, repeated=False, \n                cached=False, augmented=False):\n    dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=AUTOTUNE)\n    dataset = dataset.map(read_tfrecord, num_parallel_calls=AUTOTUNE)\n    if cached:\n        dataset = dataset.cache()\n    if shuffled:\n        dataset = dataset.shuffle(1024, seed=SEED)\n    if augmented:\n        dataset = dataset.map(data_augment, num_parallel_calls=AUTOTUNE)\n    if repeated:\n        dataset = dataset.repeat()\n    dataset = dataset.batch(BATCH_SIZE, drop_remainder=True)\n    dataset = dataset.prefetch(AUTOTUNE)\n    return dataset\n\n","execution_count":null,"outputs":[]},{"metadata":{"id":"SvcAybCiNcyD","trusted":true},"cell_type":"code","source":"# https://keras.io/examples/vision/knowledge_distillation/#construct-distiller-class\nclass Distiller(tf.keras.Model):\n    def __init__(self, student, teacher):\n        super(Distiller, self).__init__()\n        self.teacher = teacher\n        self.student = student\n\n    def compile(\n        self,\n        optimizer,\n        metrics,\n        student_loss_fn,\n        distillation_loss_fn,\n        alpha=0.1,\n    ):\n        \"\"\" Configure the distiller.\n\n        Args:\n            optimizer: Keras optimizer for the student weights\n            metrics: Keras metrics for evaluation\n            student_loss_fn: Loss function of difference between student\n                predictions and ground-truth\n            distillation_loss_fn: Loss function of difference between soft\n                student predictions and soft teacher predictions\n            alpha: weight to student_loss_fn and 1-alpha to distillation_loss_fn\n        \"\"\"\n        super(Distiller, self).compile(optimizer=optimizer, metrics=metrics)\n        self.student_loss_fn = student_loss_fn\n        self.distillation_loss_fn = distillation_loss_fn\n        self.alpha = alpha\n        \n    @tf.function\n    def train_step(self, data):\n        x, y = data\n        image, image_annotation = tf.split(x, 2, axis=1)\n        image = tf.squeeze(image)\n        image_annotation = tf.squeeze(image_annotation)\n        \n        teacher_predictions, teacher_features = self.teacher(image_annotation, training=False)\n        with tf.GradientTape() as tape:\n            student_predictions, student_features = self.student(image, training=True)\n         \n            student_loss = self.student_loss_fn(y, student_predictions)\n            student_loss = tf.reduce_sum(student_loss * (1. / BATCH_SIZE))\n            distillation_loss = self.distillation_loss_fn(tf.reshape(teacher_features, [BATCH_SIZE, -1]), tf.reshape(student_features, [BATCH_SIZE, -1]))\n            # distillation_loss = tf.reduce_sum(distillation_loss * (1. / BATCH_SIZE))\n            loss = self.alpha * student_loss + (1 - self.alpha) * distillation_loss\n            \n        # Compute gradients\n        trainable_vars = self.student.trainable_variables\n        gradients = tape.gradient(loss, trainable_vars)\n\n        # Update weights\n        self.optimizer.apply_gradients(zip(gradients, trainable_vars))\n\n        # Update the metrics configured in `compile()`.\n        self.compiled_metrics.update_state(y, student_predictions)\n\n        # Return a dict of performance\n        results = {m.name: m.result() for m in self.metrics}\n        results.update(\n            {\"student_loss\": student_loss, \"distillation_loss\": distillation_loss}\n        )\n        return results\n    \n    @tf.function\n    def valid_step(self, data):\n        x, y = data\n        image, image_annotation = tf.split(x, 2, axis=1)\n        image = tf.squeeze(image)\n        # Compute predictions\n        y_prediction, _ = self.student(image, training=False)\n\n        # Calculate the loss\n        student_loss = self.student_loss_fn(y, y_prediction)\n        student_loss = tf.reduce_sum(student_loss * (1. / BATCH_SIZE))\n        # Update the metrics.\n        self.compiled_metrics.update_state(y, y_prediction)\n\n        # Return a dict of performance\n        results = {m.name: m.result() for m in self.metrics}\n        results.update({\"student_loss\": student_loss})\n        return results\n    \n    @tf.function\n    def test_step(self, data):\n        x, y = data\n        image, image_annotation = tf.split(x, 2, axis=1)\n        image = tf.squeeze(image)\n        # Compute predictions\n        y_prediction, _ = self.student(image, training=False)\n\n        # Calculate the loss\n        student_loss = self.student_loss_fn(y, y_prediction)\n        student_loss = tf.reduce_sum(student_loss * (1. / BATCH_SIZE))\n        # Update the metrics.\n        self.compiled_metrics.update_state(y, y_prediction)\n\n        # Return a dict of performance\n        results = {m.name: m.result() for m in self.metrics}\n        results.update({\"student_loss\": student_loss})\n        return results","execution_count":null,"outputs":[]},{"metadata":{"id":"qoUr3sHuG7--","trusted":true},"cell_type":"code","source":"def get_model(name):\n    base_model = BaseModel(input_shape=(IMG_SIZE,IMG_SIZE,3), include_top=False, weights='imagenet', pooling=\"avg\")\n    base_model_output = base_model.output\n    x = L.Dropout(0.5)(base_model_output)\n    outputs = L.Dense(NUM_CLASSES, activation=\"sigmoid\")(x)\n\n    model = tf.keras.models.Model(inputs=base_model.input, outputs=[outputs, base_model_output], name=name)\n    return model\n\ndef get_distiller_model(fold=0):\n    with strategy.scope():\n        student = get_model('student')\n        teacher = get_model('teacher')\n        teacher.load_weights(f'../input/ranzcr-annotation-teacher/teacher_model_{fold}.h5')\n\n        distiller = Distiller(student=student, teacher=teacher)\n        distiller.compile(\n            optimizer=tf.keras.optimizers.Adam(lr=1e-3),\n            metrics=[tf.keras.metrics.AUC(multi_label=True)],\n            student_loss_fn=tfa.losses.SigmoidFocalCrossEntropy(alpha = 0.5, gamma = 2, reduction=tf.keras.losses.Reduction.NONE),\n            distillation_loss_fn=tf.keras.losses.MSE,\n            alpha=0.3\n        )\n    return distiller","execution_count":null,"outputs":[]},{"metadata":{"id":"e5TH_x21G-57","outputId":"6a5b3eaf-5956-46da-f527-5989db994d19","trusted":true},"cell_type":"code","source":"TF_REC_DS_PATH = KaggleDatasets().get_gcs_path('ranzcr-annotation-900-tfrecords')\n\ntfrec_files = []\nfor fold in range(5):\n    training_files = [TF_REC_DS_PATH + f'/{fold}_{num}.tfrec' for num in range(0,5)]\n    random.shuffle(training_files)\n    tfrec_files.append(training_files)","execution_count":null,"outputs":[]},{"metadata":{"id":"Zwj5TbIIHAr9","outputId":"e5b361e3-303f-405a-92ec-96d3657498d3","trusted":true},"cell_type":"code","source":"fold_item_counts = [1804, 1783, 1809, 1851, 1848]\n\nfor fold in range(5):\n    \n    if fold != target_fold:\n        continue\n        \n    print(f'fold_{fold} start')\n    train_filenames = list(itertools.chain.from_iterable([tfrec_files[i] for i in range(5) if i != fold]))\n    val_filenames = tfrec_files[fold]\n\n    random.shuffle(train_filenames)\n\n    train_dataset = get_dataset(train_filenames, shuffled=True, augmented=True, repeated=True)\n    val_dataset = get_dataset(val_filenames, shuffled=False, cached=True)\n\n    steps_per_epoch = (sum(fold_item_counts) - fold_item_counts[fold]) // BATCH_SIZE\n    validation_steps = fold_item_counts[fold] // BATCH_SIZE\n    \n    model = get_distiller_model(fold)\n    \n    sv = tf.keras.callbacks.ModelCheckpoint(f'distiller_model_{fold}.h5', monitor='val_student_loss', verbose=1, save_best_only=True, save_weights_only=True)\n    reduce_lr = tf.keras.callbacks.ReduceLROnPlateau(monitor='val_student_loss', verbose=1, factor=0.1, patience=3, min_delta=0.0001, min_lr=1e-6)\n\n    model.fit(\n        train_dataset,\n        steps_per_epoch=steps_per_epoch,\n        epochs=15,\n        callbacks=[reduce_lr, sv],\n        validation_data=val_dataset,\n    )\n    model.built = True\n    model.load_weights(f'./distiller_model_{fold}.h5')\n    model.get_layer('student').save_weights(f'student_model_{fold}.h5')\n    \n#     tf.keras.backend.clear_session()\n#     del model\n#     gc.collect()\n    ","execution_count":null,"outputs":[]}],"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}