{"cells":[{"metadata":{},"cell_type":"markdown","source":"The aim of this kernel was to create a baseline for Multi-label classification for the [Human Protein Atlas - Single Cell Classification](https://www.kaggle.com/ayuraj/hpa-multi-label-classification-with-tf-and-w-b) challenge using Tensorflow and use TPUv3-8 (refer to [this link](https://www.kaggle.com/docs/tpu) for detailed documentation of usage of TPUv3-8 on Kaggle Kernels)."},{"metadata":{},"cell_type":"markdown","source":"## Dependencies"},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install -qq visualkeras","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nimport visualkeras\nimport numpy as np\nimport pandas as pd\nfrom glob import glob\nimport tensorflow as tf\nfrom PIL import ImageFont\nfrom typing import List, Tuple\nfrom collections import Counter\nimport plotly.graph_objects as go\nfrom matplotlib import pyplot as plt\nfrom plotly.subplots import make_subplots\nfrom kaggle_datasets import KaggleDatasets","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Basic Exploratory Analysis"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Refer https://www.kaggle.com/c/hpa-single-cell-image-classification/data\n\nLABEL_NAMES = [\n    \"Nucleoplasm\", \"Nuclear Membrane\", \"Nucleoli\",\n    \"Nucleoli Fibrillar Center\", \"Nuclear Speckles\",\n    \"Nuclear Bodies\", \"Endoplasmic Reticulum\", \"Golgi Apparatus\",\n    \"Intermediate Filaments\", \"Actin Filaments\", \"Microtubules\",\n    \"Mitotic Spindle\", \"Centrosome\", \"Plasma Membrane\", \"Mitochondria\",\n    \"Aggresome\", \"Cytosol\", \"Vesicles\", \"Negative\"\n]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"GCS_PATH = KaggleDatasets().get_gcs_path('hpa-single-cell-image-classification')\n\nTRAIN_IMAGE_FILES = glob('../input/hpa-single-cell-image-classification/train/*.png')\nTEST_IMAGE_FILES = glob('../input/hpa-single-cell-image-classification/test/*.png')\n\nTRAIN_TFRECORDS = tf.io.gfile.glob(os.path.join(GCS_PATH, 'train_tfrecords/*.tfrec'))\nTEST_TFRECORDS = tf.io.gfile.glob(os.path.join(GCS_PATH, 'train_tfrecords/*.tfrec'))\n\nprint('Number of Train Images:', len(TRAIN_IMAGE_FILES))\nprint('Number of Test Images:', len(TEST_IMAGE_FILES))\nprint('Number of Train TFRecord Files:', len(TRAIN_TFRECORDS))\nprint('Number of Test TFRecord Files:', len(TEST_TFRECORDS))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataframe = pd.read_csv('../input/hpa-single-cell-image-classification/train.csv')\ndataframe.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"It is evident from the given dataframe, there are multiple labels involved with a single image, separated by `|`. Now, we would check for the distribution of classes across the dataset. We would also use a simple class weightage strategy to balance imbalanced classes."},{"metadata":{"trusted":true},"cell_type":"code","source":"# Ref: https://www.kaggle.com/dschettler8845/hpa-xai-ig-tfrecords-tpu-training\n\n# Getting Label Distributions\nLABEL_COUNTS = Counter([\n    c for sublist in dataframe['Label'].str.split('|').to_list() for c in sublist\n])\n\n# Calculating class weights\nthreshold = sorted(LABEL_COUNTS.values())[3]\nCLASS_WEIGHTS = {\n    int(k): min(1.0, threshold / v) for k, v in LABEL_COUNTS.items()\n}\n\n# Visualization of Label imbalance and class weights\nfig = go.Figure(data=[\n    go.Bar(\n        name='Class Distributions',\n        x=[str(i) for i in list(range(len(LABEL_NAMES)))],\n        y=[LABEL_COUNTS[str(key)] for key in list(range(len(LABEL_NAMES)))]\n    ),\n    go.Bar(\n        name='Class Weights',\n        x=[str(i) for i in list(range(len(LABEL_NAMES)))],\n        y=[CLASS_WEIGHTS[key] * 2000 for key in list(CLASS_WEIGHTS.keys())]\n    )\n])\nfig.update_layout(barmode='stack', uniformtext_minsize=8, uniformtext_mode='hide')\nfig.update_traces(textposition='outside')\nfig.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Tensorflow Dataloader from TFRecords"},{"metadata":{},"cell_type":"markdown","source":"Since TPUs are very fast, many models ported to TPU end up with a data bottleneck. The TPU is sitting idle, waiting for data for the most part of each training epoch. TPUs read training data exclusively from GCS buckets. Data for TPU training typically comes sharded across the appropriate number of larger TFRecord files. We would create a TFRecord Dataloader class to read the data from tfrecord files."},{"metadata":{"trusted":true},"cell_type":"code","source":"class TFRecordLoader:\n\n    def __init__(self, image_size: List[int], n_classes: int, include_yellow_channel: bool):\n        self.image_size = image_size\n        self.n_classes = n_classes\n        self.include_yellow_channel = include_yellow_channel\n\n    def _parse_image(self, image):\n        image = tf.image.decode_png(image, channels=1)\n        image = tf.image.resize(image, self.image_size)\n        image = tf.cast(image, dtype=tf.float32)\n        return image\n\n    def _parse_label(self, label):\n        indices = tf.strings.to_number(\n            tf.strings.split(label, sep='|'),\n            out_type=tf.int32\n        )\n        return tf.reduce_max(\n            tf.one_hot(indices, depth=self.n_classes), axis=-2\n        )\n\n    def _make_example(self, example):\n        feature_format = {\n            'image': tf.io.FixedLenFeature(shape=(), dtype=tf.string),\n            'image_name': tf.io.FixedLenFeature(shape=(), dtype=tf.string),\n            'target': tf.io.FixedLenFeature(shape=(), dtype=tf.string)\n        }\n        features = tf.io.parse_single_example(example, features=feature_format)\n        image = self._parse_image(features['image'])\n        image_name = features['image_name']\n        label = self._parse_label(features['target'])\n        return image, image_name, label\n\n    def _combine_channels(self, red, green, blue, yellow):\n        # Ref: https://www.kaggle.com/dschettler8845/hpa-xai-ig-tfrecords-tpu-training\n        (r_i, r_j, r_k), (g_i, g_j, g_k), (b_i, b_j, b_k), (y_i, y_j, y_k) = red, green, blue, yellow\n        combined_image = tf.stack(\n            [r_i[..., 0], g_i[..., 0], b_i[..., 0], y_i[..., 0]], axis=-1\n        ) if self.include_yellow_channel else tf.stack(\n            [r_i[..., 0], g_i[..., 0], b_i[..., 0]], axis=-1\n        )\n        return combined_image, r_k\n\n    def _preprocess(self, dataset):\n        # Ref: https://www.kaggle.com/dschettler8845/hpa-xai-ig-tfrecords-tpu-training\n        red_dataset = dataset.filter(\n            lambda x, y, z: tf.strings.regex_full_match(y, \".*red.*\"))\n        green_dataset = dataset.filter(\n            lambda x, y, z: tf.strings.regex_full_match(y, \".*green.*\"))\n        blue_dataset = dataset.filter(\n            lambda x, y, z: tf.strings.regex_full_match(y, \".*blue.*\"))\n        yellow_dataset = dataset.filter(\n            lambda x, y, z: tf.strings.regex_full_match(y, \".*yellow.*\"))\n        dataset = tf.data.Dataset.zip(\n            (red_dataset, green_dataset, blue_dataset, yellow_dataset)\n        )\n        dataset = dataset.map(\n            map_func=self._combine_channels,\n            num_parallel_calls=tf.data.AUTOTUNE\n        )\n        return dataset\n\n    def get_dataset(self, train_tfrecord_files: List[str], ignore_order: bool = False):\n        options = tf.data.Options()\n        options.experimental_deterministic = False\n        dataset = tf.data.TFRecordDataset(\n            train_tfrecord_files, num_parallel_reads=tf.data.AUTOTUNE)\n        dataset = dataset.with_options(options) if ignore_order else dataset\n        dataset = dataset.map(\n            map_func=self._make_example, num_parallel_calls=tf.data.AUTOTUNE)\n        dataset = self._preprocess(dataset)\n        return dataset","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def plot_result(\n    images, captions: List[str], title: str, figsize: Tuple[int, int]):\n    fig = plt.figure(figsize=figsize)\n    plt.suptitle(\n        'Label: ' + title[0], fontsize=20, fontweight='bold')\n    for index in range(len(captions)):\n        fig.add_subplot(\n            1, len(captions), index + 1\n        ).set_title(captions[index])\n        _ = plt.imshow(images[index])\n        plt.axis(False)\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"loader = TFRecordLoader(\n    image_size=[512, 512], n_classes=19, include_yellow_channel=False\n)\ndataset = loader.get_dataset(TRAIN_TFRECORDS)\n\nfor x, y in dataset.take(4):\n    plot_result(\n        [x[..., 0], x[..., 1], x[..., 2], x],\n        ['red channel', 'green channel', 'blue channel', 'combined image'],\n        [LABEL_NAMES[label] for label in np.where(y.numpy()==1)[0]], (20, 6)\n    )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"loader = TFRecordLoader(\n    image_size=[512, 512], n_classes=19, include_yellow_channel=True\n)\ndataset = loader.get_dataset(TRAIN_TFRECORDS)\n\nfor x, y in dataset.take(4):\n    plot_result(\n        [x[..., 0], x[..., 1], x[..., 2], x[..., 3]],\n        ['red channel', 'green channel', 'blue channel', 'yellow channel'],\n        [LABEL_NAMES[label] for label in np.where(y.numpy()==1)[0]], (20, 6)\n    )","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Data Augmentation"},{"metadata":{"trusted":true},"cell_type":"code","source":"class AugmentationFactory:\n\n    def __init__(self, include_flips: bool, include_rotation: bool, include_jitter: bool):\n        self.include_flips = include_flips\n        self.include_rotation = include_rotation\n        self.include_jitter = include_jitter\n\n    @staticmethod\n    def _flip_horizontal(image, seed):\n        image = tf.image.stateless_random_flip_left_right(image, seed)\n        return image\n\n    @staticmethod\n    def _flip_vertical(image, seed):\n        image = tf.image.stateless_random_flip_up_down(image, seed)\n        return image\n\n    @staticmethod\n    def _rotate(image):\n        rotation_k = tf.random.uniform((1,), minval=0, maxval=4, dtype=tf.int32)[0]\n        image = tf.image.rot90(image, k=rotation_k)\n        return image\n\n    @staticmethod\n    def _random_jitter(image, seed):\n        image = tf.image.stateless_random_saturation(image, 0.9, 1.1, seed)\n        image = tf.image.stateless_random_brightness(image, 0.075, seed)\n        image = tf.image.stateless_random_contrast(image, 0.9, 1.1, seed)\n        return image\n\n    def _map_augmentations(self, image, label):\n        seed = tf.random.uniform((2,), minval=0, maxval=100, dtype=tf.int32)\n        if self.include_flips:\n            image = self._flip_horizontal(image=image, seed=seed)\n            image = self._flip_vertical(image=image, seed=seed)\n        image = self._rotate(image=image) if self.include_rotation else image\n        image = self._random_jitter(image=image, seed=seed) if self.include_jitter else image\n        return image, label\n\n    def augment_dataset(self, dataset):\n        return dataset.map(\n            map_func=self._map_augmentations,\n            num_parallel_calls=tf.data.AUTOTUNE\n        )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"loader = TFRecordLoader(\n    image_size=[512, 512], n_classes=19, include_yellow_channel=False\n)\ndataset = loader.get_dataset(TRAIN_TFRECORDS)\n\naugmentation_factory = AugmentationFactory(\n    include_flips=True, include_rotation=True, include_jitter=True\n)\ndataset = augmentation_factory.augment_dataset(dataset)\n\nfor x, y in dataset.take(4):\n    plot_result(\n        [x[..., 0], x[..., 1], x[..., 2], x],\n        ['red channel', 'green channel', 'blue channel', 'Combined Image'],\n        [LABEL_NAMES[label] for label in np.where(y.numpy()==1)[0]], (20, 6)\n    )","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Simple Tensorflow Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_backbone(backbone_name: str, input_shape: List[int], weights: str = 'imagenet'):\n    backbone_class = None\n    if 'efficientnet' in backbone_name:\n        if 'b0' in backbone_name:\n            backbone_class = tf.keras.applications.EfficientNetB0\n        elif 'b1' in backbone_name:\n            backbone_class = tf.keras.applications.EfficientNetB1\n        elif 'b2' in backbone_name:\n            backbone_class = tf.keras.applications.EfficientNetB2\n        elif 'b3' in backbone_name:\n            backbone_class = tf.keras.applications.EfficientNetB3\n        elif 'b4' in backbone_name:\n            backbone_class = tf.keras.applications.EfficientNetB4\n        elif 'b5' in backbone_name:\n            backbone_class = tf.keras.applications.EfficientNetB5\n        elif 'b6' in backbone_name:\n            backbone_class = tf.keras.applications.EfficientNetB6\n        elif 'b7' in backbone_name:\n            backbone_class = tf.keras.applications.EfficientNetB7\n    backbone = backbone_class(\n        include_top=False, weights=weights, input_shape=input_shape\n    )\n    return backbone\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Ref: https://www.kaggle.com/ayuraj/hpa-multi-label-classification-with-tf-and-w-b\n\ndef simple_model(input_shape: List[int], backbone_name: str, dropout: float, n_classes: int):\n    backbone = get_backbone(\n        backbone_name=backbone_name, weights='imagenet', input_shape=input_shape\n    )\n    backbone.trainable = True\n    input_tensor = tf.keras.Input(input_shape, name='inputs')\n    backbone_features = backbone(input_tensor, training=True)\n    x = tf.keras.layers.GlobalAveragePooling2D(name='global_average_pool_2d')(backbone_features)\n    x = tf.keras.layers.Dropout(rate=dropout, name='dropout_{}'.format(dropout))(x) if dropout > 0 else x\n    output_tensor = tf.keras.layers.Dense(n_classes, activation='softmax', name='outputs')(x)\n    return tf.keras.Model(\n        input_tensor, output_tensor,\n        name='{}_transfer_learning_model'.format(backbone_name)\n    )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_strategy():\n    try:  # detect TPUs\n        tpu = tf.distribute.cluster_resolver.TPUClusterResolver()  # TPU detection\n        tf.config.experimental_connect_to_cluster(tpu)\n        tf.tpu.experimental.initialize_tpu_system(tpu)\n        strategy = tf.distribute.experimental.TPUStrategy(tpu)\n    except ValueError:  # detect GPUs\n        strategy = tf.distribute.MirroredStrategy()  # for GPU or multi-GPU machines\n    print(\"Number of accelerators: \", strategy.num_replicas_in_sync)\n    return strategy","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"strategy = get_strategy()\n\nwith strategy.scope():\n    model = simple_model(\n        input_shape=[224, 224, 3], backbone_name='efficientnetb0',\n        dropout=0.5, n_classes=len(LABEL_NAMES)\n    )\nmodel.compile(\n    optimizer=tf.keras.optimizers.Adam(),\n    loss=tf.keras.losses.BinaryCrossentropy(label_smoothing=0.0),\n    metrics=[tf.keras.metrics.AUC(multi_label=True)]\n)\nmodel.summary()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"visualkeras.layered_view(model, spacing=100)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"For the baseline, I created a simple model to perform transfer learning from EfficientNet Backbones pre-trained on Imagenet. Unlike the baseline kernel by [Ayush Thakur](https://www.kaggle.com/ayuraj), where Sigmoid Focal Crossentrpy Loss was used, I used plain old Binary Crossentropy. While using Focal Loss, it seems to be converging faster and in order to acheive similar results, I tried experimenting with a custor learning rate scheduling strategy which we will discuss shortly."},{"metadata":{},"cell_type":"markdown","source":"## Training"},{"metadata":{},"cell_type":"markdown","source":"### Configure Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"def configure_dataset(augmented_dataset, shuffle_buffer: int = 128, batch_size: int = 16):\n    dataset = augmented_dataset.repeat()\n    dataset = augmented_dataset.shuffle(shuffle_buffer)\n    dataset = dataset.batch(batch_size)\n    dataset = dataset.prefetch(buffer_size=tf.data.AUTOTUNE)\n    return dataset","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"loader = TFRecordLoader(\n    image_size=[224, 224], n_classes=19,\n    include_yellow_channel=False\n)\ndataset = loader.get_dataset(\n    TRAIN_TFRECORDS, ignore_order=True\n)\n\naugmentation_factory = AugmentationFactory(\n    include_flips=True, include_rotation=True, include_jitter=True\n)\ndataset = augmentation_factory.augment_dataset(dataset)\nBATCH_SIZE = 64 * strategy.num_replicas_in_sync\ntrain_dataset = configure_dataset(\n    dataset, batch_size=BATCH_SIZE\n)\ntrain_dataset","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Custom Learning Rate Scheduling"},{"metadata":{"trusted":true},"cell_type":"code","source":"def custom_lr_scheduler(epoch, warmup_epochs=3, sustain_epochs=2, initial_lr=1e-5, max_lr=1e-4, epsilon=0.9):\n    if epoch < warmup_epochs:\n        lr = ((max_lr - initial_lr) / warmup_epochs * epoch) + initial_lr\n    elif epoch < warmup_epochs + sustain_epochs:\n        lr = max_lr\n    else:\n        lr = ((max_lr - initial_lr) * epsilon ** (epoch - warmup_epochs)) + initial_lr\n    return lr","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The idea of the custom Learning rate scheduling strategy was inspired by the kernel [HPA - XAI & IG [TFRECORDS+TPU][TRAINING]](https://www.kaggle.com/dschettler8845/hpa-xai-ig-tfrecords-tpu-training) by [Darien Schettler](https://www.kaggle.com/dschettler8845). The idea is basically to increase the learning rate at a constant rate initially for faster convergence and then keep decreasing it exponentially post the warmup epochs. The initial increase and subsequent sutainance of the increased learning rate helps the model to converge faster while the eponential decay of learning rate susequently helps in avoiding overfitting, especially while training for a large number of epochs"},{"metadata":{"trusted":true},"cell_type":"code","source":"fig = go.Figure(\n    data=go.Scatter(\n        x=list(range(1, 16)),\n        y=[\n            custom_lr_scheduler(epoch) for epoch in list(range(1, 16))\n        ]\n    )\n)\nfig.update_layout(title='Custom LR Scheduling Policy')\nfig.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Training the Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"callbacks = [\n    tf.keras.callbacks.LearningRateScheduler(custom_lr_scheduler),\n    tf.keras.callbacks.EarlyStopping(monitor='loss', patience=10, restore_best_weights=True),\n    tf.keras.callbacks.ReduceLROnPlateau(monitor='loss', patience=5)\n]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"EPOCHS = 150\n\nhistory = model.fit(\n    train_dataset, epochs=EPOCHS,\n    class_weight=CLASS_WEIGHTS\n)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Save Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"model.save(\n    './weights/efficientnetb0_fine_tuned_classification',\n    options=tf.saved_model.SaveOptions(\n        experimental_io_device='/job:localhost'\n    )\n)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Training History"},{"metadata":{"trusted":true},"cell_type":"code","source":"fig = make_subplots(rows=2, cols=1)\nfig.add_trace(\n    go.Scatter(\n        x=list(range(EPOCHS)),\n        y=history.history['loss'],\n         name='loss'\n    ), row=1, col=1\n)\nfig.add_trace(\n    go.Scatter(\n        x=list(range(EPOCHS)),\n        y=history.history['auc'],\n         name='AUC'\n    ), row=2, col=1\n)\nfig.update_layout(title='Training History')\nfig.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Evaluating the Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"loader = TFRecordLoader(\n    image_size=[224, 224], n_classes=19, include_yellow_channel=False\n)\ndataset = loader.get_dataset(TEST_TFRECORDS)\ntest_dataset = configure_dataset(\n    dataset, batch_size=BATCH_SIZE\n)\n\ntest_loss, test_auc = model.evaluate(test_dataset)\n\nprint('Loss on Test Data:', test_loss)\nprint('Loss on Test AUC:', test_auc)","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}