{"cells":[{"metadata":{},"cell_type":"markdown","source":"# TPU Training Code for EfficientNet B4\n* Sigmoid Focal Cross Entropy Loss with Label Smoothing 0.1\n* Learning Rate Sheduler with Warmup and Cosine Decay\n* Light augmentations (Flipping, Saturation...) - No CutOut/CutMix\n* 20 epochs CV (using predefined folds so that it was easier to ensemble different models). Callback for saving weights of best epoch\n* Final model was trained using best avg. number of epochs in CV phase (in this case 14)\n* TF Records from 2020 competition dataset using the 512x512px images. Smaller resolutions have tended to degrade the model.\n \n\nSpecial thanks to [Mark Wijkhuizen](https://www.kaggle.com/markwijkhuizen/tf-efficientnetb4-mixup-cutmix-gridmask-cv-0-90) from whose notebook the learning rate timer originates and to [ragner123](https://www.kaggle.com/ragnar123/effb5-cv-0-9007-single-model-tf) from whom I borrowed some of the TPU image augmentations (including parameters)."},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nimport json\nimport sys\nimport random\nimport pandas as pd\nimport numpy as np\nimport tensorflow as tf\nimport tensorflow_probability as tfp\nimport itertools\nimport joblib\nimport math\nfrom tensorflow import keras\nfrom keras.preprocessing.image import ImageDataGenerator\nfrom keras.models import load_model\nfrom keras.utils import plot_model\nfrom tensorflow.keras.experimental import CosineDecay\nfrom tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping\nfrom tensorflow.keras import backend as K\nfrom PIL import Image\nfrom functools import partial\nfrom kaggle_datasets import KaggleDatasets","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"tf.__version__","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"try:\n    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n    print('Device:', tpu.master())\n    tf.config.experimental_connect_to_cluster(tpu)\n    tf.tpu.experimental.initialize_tpu_system(tpu)\n    strategy = tf.distribute.experimental.TPUStrategy(tpu)\n    DEVICE = \"TPU\"\nexcept:\n    DEVICE = \"notTPU\"\n    strategy = tf.distribute.get_strategy()\n\nAUTOTUNE = tf.data.experimental.AUTOTUNE\nGCS_PATH = KaggleDatasets().get_gcs_path(f'cassava-leaf-disease-classification')\nREPLICAS =  strategy.num_replicas_in_sync\nBATCH_SIZE = 32 * strategy.num_replicas_in_sync\nFILENAMES = tf.io.gfile.glob(GCS_PATH + '/train_tfrecords/*.tfrec')\n\nLR = 2e-6 * REPLICAS\nEPOCHS = 20\n\nimg_size = (512, 512)\ndefault_img_size = (512, 512) ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Loading data and use light image augmentations"},{"metadata":{"trusted":true},"cell_type":"code","source":"def decode_image(image, img_size):\n    image = tf.image.decode_jpeg(image, channels=3)\n    if img_size != default_img_size:\n        image = tf.image.resize(image, img_size)\n    image = tf.cast(image, tf.bfloat16)\n    return tf.reshape(image, [*img_size, 3])\n\ndef _parse_image(proto, train):\n    features = {\"image_name\": tf.io.FixedLenFeature([], tf.string),\n                \"image\": tf.io.FixedLenFeature([], tf.string)}\n    if train:\n        features[\"target\"] = tf.io.FixedLenFeature([], tf.int64)\n    return tf.io.parse_single_example(proto, features)\n\ndef get_image_and_label(proto, train, img_size):\n    sample = _parse_image(proto, train=train)\n    sample[\"image\"] = decode_image(sample[\"image\"], img_size)\n    if train:\n        return sample[\"image\"], tf.one_hot(tf.cast(sample['target'], tf.int32), 5)\n    return sample[\"image\"], sample[\"image_name\"]\n\n\ndef image_augmentations(image,label):\n    \n    image = tf.image.random_flip_left_right(image)\n    image = tf.image.random_flip_up_down(image)\n    \n    if tf.random.uniform([], 0, 1.0, dtype = tf.float32) > 0.75:\n        image = tf.image.transpose(image)\n    \n    probablity_rotation = tf.random.uniform([], 0, 1.0, dtype = tf.float32)\n    if probablity_rotation > 0.75:\n        image = tf.image.rot90(image, k = 3)\n    elif probablity_rotation > 0.5:\n        image = tf.image.rot90(image, k = 2)\n    elif probablity_rotation > 0.25:\n        image = tf.image.rot90(image, k = 1)\n        \n    if tf.random.uniform([], 0, 1.0, dtype = tf.float32) >= 0.4:\n        image = tf.image.random_saturation(image, lower = 0.8, upper = 1.2)\n    if tf.random.uniform([], 0, 1.0, dtype = tf.float32) >= 0.4:\n        image = tf.image.random_contrast(image, lower = 0.8, upper = 1.2)\n    if tf.random.uniform([], 0, 1.0, dtype = tf.float32) >= 0.4:\n        image = tf.image.random_brightness(image, max_delta = 0.1)\n    \n    probability_cropping = tf.random.uniform([], 0, 1.0, dtype = tf.float32)\n    if probability_cropping > 0.7:\n        if probability_cropping > 0.9:\n            image = tf.image.central_crop(image, central_fraction = 0.7)\n        elif probability_cropping > 0.8:\n            image = tf.image.central_crop(image, central_fraction = 0.8)\n        else:\n            image = tf.image.central_crop(image, central_fraction = 0.9)\n    elif probability_cropping > 0.5:\n        crop_size = tf.random.uniform([], int(img_size[0] * 0.8), img_size[0], dtype = tf.int32)\n        image = tf.image.random_crop(image, size = [crop_size, crop_size, 3])\n\n    image = tf.image.resize(image, size = img_size)\n    image = tf.reshape(image, [*img_size, 3])\n    \n    return image, label\n\n\ndef get_training_dataset(tfr, batchsize, img_size=default_img_size):\n    return tfr.map(partial(get_image_and_label, train=True, img_size=img_size)).repeat().map(\n        lambda x,y: image_augmentations(x,y)).shuffle(1000).batch(batchsize).prefetch(AUTOTUNE)\n\ndef get_validation_dataset(tfr, batchsize, img_size=default_img_size, train=True):\n    return tfr.map(partial(get_image_and_label, train=train, img_size=img_size)).batch(batchsize).prefetch(AUTOTUNE)\n\n# Scale normalization batch as normalization layer is after rescaling layer in b4\ndef get_normalization_batch(tfr, batchsize, img_size=default_img_size):\n    return tfr.map(partial(get_image_and_label, train=True, img_size=img_size)).map(\n        lambda x,y: x/255.0).shuffle(1000).batch(batchsize).prefetch(AUTOTUNE)\n\ndef get_tfrecord_size(tfrecord):\n    return sum(1 for _ in tfrecord)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# For infos on LR scheduler please refer to https://www.kaggle.com/markwijkhuizen/tf-efficientnetb4-mixup-cutmix-gridmask-cv-0-90)\ndef lrfn(epoch, bs=BATCH_SIZE, epochs=20):\n\n    LR_START = 1e-6\n    LR_MAX = 2e-4\n    LR_FINAL = 1e-6\n    LR_RAMPUP_EPOCHS = 4\n    LR_SUSTAIN_EPOCHS = 0\n    DECAY_EPOCHS = epochs  - LR_RAMPUP_EPOCHS - LR_SUSTAIN_EPOCHS - 1\n    LR_EXP_DECAY = (LR_FINAL / LR_MAX) ** (1 / (EPOCHS - LR_RAMPUP_EPOCHS - LR_SUSTAIN_EPOCHS - 1))\n\n    if epoch < LR_RAMPUP_EPOCHS:\n        lr = LR_START + (LR_MAX + LR_START) * (epoch / LR_RAMPUP_EPOCHS) ** 2.5\n    elif epoch < LR_RAMPUP_EPOCHS + LR_SUSTAIN_EPOCHS:\n        lr = LR_MAX\n    else:\n        epoch_diff = epoch - LR_RAMPUP_EPOCHS - LR_SUSTAIN_EPOCHS\n        decay_factor = (epoch_diff / DECAY_EPOCHS) * math.pi\n        decay_factor= (tf.math.cos(decay_factor).numpy() + 1) / 2        \n        lr = LR_FINAL + (LR_MAX - LR_FINAL) * decay_factor\n\n    return lr\n\nimport matplotlib.pyplot as plt  \nrng = [i for i in range(EPOCHS)]\ny = [lrfn(x) for x in rng]\nplt.plot(rng, y);","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Efficientnet B4 with NoisyStudent weights\nEfficientNetB4 is finetuned based on pretrained NoisyStudent weights. Weights can e.g. be downloaded from the official[ Googleapi site](https://storage.googleapis.com/cloud-tpu-checkpoints/efficientnet/noisystudent/noisy_student_efficientnet-b4.tar.gz).\nThe checkpoint files can be transformed to an .h5 file using the code from the [Tensorflow Github page ](https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet).\n\nIn the Cassava Classification Challenge, the use of NoisyStudent only led to a very small increase in performance.\n"},{"metadata":{"trusted":true},"cell_type":"code","source":"from tensorflow.keras.applications import EfficientNetB4\n\ndef efficientnetb4(weights=\"../input/cassava-leaf-disease-efficientnetb4/efficientnetb4_noisystudent_notop.h5/efficientnetb4_noisystudent_notop.h5\",\n                   dropout=(0.4, 0.5)):\n    \n    keras.backend.reset_uids()\n    \n    efficientnet = EfficientNetB4(weights=weights, include_top=False,\n                                  input_shape=(*img_size, 3), drop_connect_rate=dropout[0])\n    \n    for layer in reversed(efficientnet.layers):\n        if isinstance(layer, tf.keras.layers.BatchNormalization):\n            layer.trainable = False\n        else:\n            layer.trainable = True\n    \n    inputs = keras.layers.Input(shape=(*img_size, 3))\n    \n    efficientnet = efficientnet(inputs)\n    pooling = keras.layers.GlobalAveragePooling2D()(efficientnet)\n    dropout = keras.layers.Dropout(dropout[1])(pooling)\n    outputs = keras.layers.Dense(5, activation=\"softmax\")(dropout)\n    \n    model = keras.Model(inputs=inputs, outputs=outputs)\n\n    return model","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We did experiments with different approaches to handle the imbalanced and noisy data. It turned among all tested loss functions, sigmoid focal crossentropy worked best especially if combined with label smoothing (factor 0.1).\n\nYou can find more info on the loss on the [TF Addon page ](https://www.tensorflow.org/addons/api_docs/python/tfa/losses/SigmoidFocalCrossEntropy)"},{"metadata":{"trusted":true},"cell_type":"code","source":"#@tf.function\ndef sigmoid_focal_crossentropy(y_true, y_pred, alpha=0.25, gamma=2.0):\n    \n    def smooth(y, smooth_factor):\n        assert len(y.shape) == 2\n        y *= 1 - smooth_factor\n        y += smooth_factor / y.shape[1]\n        return y\n    \n    #label smoothing factor\n    FACTOR = 0.1\n    alpha_factor = 1.0\n    modulating_factor = 1.0\n\n    y_pred = tf.convert_to_tensor(y_pred)\n    y_true = tf.convert_to_tensor(smooth(y_true, FACTOR), dtype=y_pred.dtype)\n\n    ce = K.binary_crossentropy(y_true, y_pred, from_logits=False)\n\n    p_t = (y_true * y_pred) + ((1 - y_true) * (1 - y_pred))\n    \n    alpha = tf.convert_to_tensor(alpha, dtype=K.floatx())\n    alpha_factor = y_true * alpha + (1 - y_true) * (1 - alpha)\n\n    gamma = tf.convert_to_tensor(gamma, dtype=K.floatx())\n    modulating_factor = tf.pow((1.0 - p_t), gamma)\n\n    return tf.reduce_sum(alpha_factor * modulating_factor * ce, axis=-1)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 5-Fold CV Training for finding best number of epochs\nWe used the same CV splits for almost all models to make it easier to build and evaluate ensembles later with the OOF predictions.<br>\n<code>\nkfold = KFold(n_splits=5, shuffle=True, random_state=123)\nfilenames = [train_path + f for f in os.listdir(train_path)]\nfolds = {}\nfor fold, (train_idx, val_idx) in enumerate(kfold.split(filenames)):\n  folds[fold] = {\"split\": (np.take(filenames, train_idx), np.take(filenames, val_idx))}\n</code>\n<br>\nSince the training lasted several hours, we performed the training for each fold individually."},{"metadata":{"trusted":true},"cell_type":"code","source":"ignore_order = tf.data.Options()\nignore_order.experimental_deterministic = False\nfolds = joblib.load(\"../input/cassava-leaf-disease-efficientnetb4/folds.pkl\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"cv_history = []\n\n#Considered folds of current iteration\ncurrent_folds = [4]\n\nfor fold, settings in dict(map(lambda k: (k, folds[k]), current_folds)).items():\n\n  tf.tpu.experimental.initialize_tpu_system(tpu)\n\n  validation_records = tf.data.TFRecordDataset(\n      [\"gs://\"+FILENAMES[0].split(\"/\")[2]+\"/train_tfrecords/\" + f.split(\"/\")[-1] for f in settings[\"split\"][1]], num_parallel_reads=AUTOTUNE)\n    \n  train_records = tf.data.TFRecordDataset(\n      [\"gs://\"+FILENAMES[0].split(\"/\")[2]+\"/train_tfrecords/\" + f.split(\"/\")[-1] for f in settings[\"split\"][0]], num_parallel_reads=AUTOTUNE)\n  train_records = train_records.with_options(ignore_order)\n\n  validation_data = get_validation_dataset(validation_records, BATCH_SIZE, img_size, train=True) \n  train_data = get_training_dataset(train_records, BATCH_SIZE, img_size)\n  \n  train_size = get_tfrecord_size(train_records)\n  validation_size = get_tfrecord_size(validation_records)\n\n  with strategy.scope():\n    \n    model = efficientnetb4()\n    # Adapt normalization layer to mean and std of dataset\n    model.get_layer('efficientnetb4').get_layer('normalization').adapt(get_normalization_batch(train_records, BATCH_SIZE, img_size))\n    \n    opt = keras.optimizers.Adam(lr=LR)\n    model.compile(loss=sigmoid_focal_crossentropy, optimizer=opt, metrics=[\"categorical_accuracy\"])\n\n  model.summary()\n    \n  cb_checkpoint = tf.keras.callbacks.ModelCheckpoint(\"backup.h5\",monitor=\"val_categorical_accuracy\", verbose=1, save_best_only=True)\n  cb_earlystop = keras.callbacks.EarlyStopping(monitor='val_categorical_accuracy', mode='max', \n                       patience=4, restore_best_weights=True, verbose=1)\n\n  cb_lr = tf.keras.callbacks.LearningRateScheduler(lambda epoch: lrfn(epoch), verbose=1)\n\n  params = {\"epochs\":EPOCHS, \"steps_per_epoch\":train_size//BATCH_SIZE,\n            \"validation_data\": validation_data,\n            \"callbacks\": [cb_earlystop, cb_checkpoint, cb_lr]} \n  \n  print(f\"Fold: {fold}, {train_size} train images {validation_size} validation images\")\n\n  history = model.fit(train_data, **params)\n    \n  cv_history.append(history.history)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training on all data with 14 epochs\nThe model for the final submission was trained with 14 epochs. We only increased the best average epoch of the CV folds by ~25% because we have more data when training without a validation set."},{"metadata":{"trusted":true},"cell_type":"code","source":"# tf.tpu.experimental.initialize_tpu_system(tpu)\n\n# train_records = tf.data.TFRecordDataset(FILENAMES, num_parallel_reads=AUTOTUNE)\n# train_records = train_records.with_options(ignore_order)\n\n# train_data = get_training_dataset(train_records, BATCH_SIZE, img_size)\n# train_size = get_tfrecord_size(train_records)\n\n# with strategy.scope():\n#     model = efficientnetb4()\n#     model.get_layer('efficientnetb4').get_layer('normalization').adapt(get_normalization_batch(trec, BATCH_SIZE, img_size))\n#     opt = keras.optimizers.Adam(lr=LR)\n#     model.compile(loss=sigmoid_focal_crossentropy, optimizer=opt, metrics=[\"categorical_accuracy\"])\n    \n# model.summary()\n\n# lr_callback = tf.keras.callbacks.LearningRateScheduler(lambda epoch: lrfn(epoch), verbose=1)\n# params = {\"epochs\": 14, \"steps_per_epoch\":train_size//BATCH_SIZE, \"callbacks\": [lr_callback]} \n  \n# print(f\"Training on whole data using {train_size} images\")\n# history = model.fit(train_data, **params)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# model.save(\"B4_14eps.h5\")","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}