{"cells":[{"metadata":{},"cell_type":"markdown","source":"### Big Transfer (BiT) pretrained models\nThis notebook is a demonstration on how to use BiT models for transfer learning in tensorflow 2.0. Since there is still problems with using tf.hub with kaggle notebooks, I uploaded some of the models on kaggle as a [dataset](https://www.kaggle.com/ibrahimsherify/bit-tf-hub-models) to be able to use them. \n\nBiT models were specifically designed for transfer learning. The authors of the paper also provided the best techniques for achieving best results. This notebook was inspired heavily by the BiT blog published on tensorflow blog outlining how to use the model and deriving the best fine tuning technique.\n\nThe main fine tuning recepies to choose from depends on the dataset according to the paper.\n- **Image size 512x512:** Apply a random crop size of 480x480 and random horizontal flip \n- **Dataset size less than 20k:** schedule length of 500 steps and no mix-up\n- **BiT-HyperRule for lr:** Decay learning rate by a factor of 10 at 3 different boundry steps\n\nReferences: <br>\nhttps://colab.research.google.com/github/google-research/big_transfer/blob/master/colabs/big_transfer_tf2.ipynb <br>\nhttps://blog.tensorflow.org/2020/05/bigtransfer-bit-state-of-art-transfer-learning-computer-vision.html<br>\nhttps://arxiv.org/abs/1912.11370"},{"metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","trusted":true},"cell_type":"code","source":"import os\nimport re\nimport math\nimport warnings\nimport numpy as np\nimport pandas as pd\n\nimport seaborn as sns\nfrom matplotlib import pyplot as plt\nfrom kaggle_datasets import KaggleDatasets\n\nfrom sklearn.metrics import f1_score, precision_score, recall_score, confusion_matrix\n\nimport tensorflow as tf\nimport tensorflow_hub as hub\nimport tensorflow.keras.backend as K\nprint(tf.__version__)","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_kg_hide-input":true,"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","trusted":true},"cell_type":"code","source":"def seed_everything(seed=0):\n    np.random.seed(seed)\n    tf.random.set_seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    os.environ['TF_DETERMINISTIC_OPS'] = '1'\n    \nsns.set()\nSEED = 17\nseed_everything(SEED)\nwarnings.filterwarnings(\"ignore\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Configurations"},{"metadata":{"trusted":true},"cell_type":"code","source":"# TPU or GPU detection\n# Detect hardware, return appropriate distribution strategy\ntry:\n    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n    print('Running on TPU ', tpu.master())\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\nprint(\"REPLICAS: \", strategy.num_replicas_in_sync)","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"HEIGHT, WIDTH = (512, 512)\nCHANNELS = 3\nN_CLASSES = 104\n\n#Get datatset from google cloud services\nGCS_PATH = KaggleDatasets().get_gcs_path('tpu-getting-started') + '/tfrecords-jpeg-%sx%s' % (HEIGHT, WIDTH)\n\nTRAINING_FILENAMES = tf.io.gfile.glob(GCS_PATH + '/train/*.tfrec')\nVALIDATION_FILENAMES = tf.io.gfile.glob(GCS_PATH + '/val/*.tfrec')\nTEST_FILENAMES = tf.io.gfile.glob(GCS_PATH + '/test/*.tfrec')\n\nCLASSES = [\n    'pink primrose', 'hard-leaved pocket orchid', 'canterbury bells', 'sweet pea', \n    'wild geranium', 'tiger lily', 'moon orchid', 'bird of paradise', 'monkshood', \n    'globe thistle', 'snapdragon', \"colt's foot\", 'king protea', 'spear thistle', \n    'yellow iris', 'globe-flower', 'purple coneflower', 'peruvian lily', \n    'balloon flower', 'giant white arum lily', 'fire lily', 'pincushion flower', \n    'fritillary', 'red ginger', 'grape hyacinth', 'corn poppy', \n    'prince of wales feathers', 'stemless gentian', 'artichoke', 'sweet william', \n    'carnation', 'garden phlox', 'love in the mist', 'cosmos',  'alpine sea holly', \n    'ruby-lipped cattleya', 'cape flower', 'great masterwort',  'siam tulip', \n    'lenten rose', 'barberton daisy', 'daffodil',  'sword lily', 'poinsettia', \n    'bolero deep blue',  'wallflower', 'marigold', 'buttercup', 'daisy', \n    'common dandelion', 'petunia', 'wild pansy', 'primula',  'sunflower', \n    'lilac hibiscus', 'bishop of llandaff', 'gaura',  'geranium', 'orange dahlia', \n    'pink-yellow dahlia', 'cautleya spicata',  'japanese anemone', 'black-eyed susan', \n    'silverbush', 'californian poppy',  'osteospermum', 'spring crocus', 'iris', \n    'windflower',  'tree poppy', 'gazania', 'azalea', 'water lily',  'rose', \n    'thorn apple', 'morning glory', 'passion flower',  'lotus', 'toad lily', \n    'anthurium', 'frangipani',  'clematis', 'hibiscus', 'columbine', 'desert-rose', \n    'tree mallow', 'magnolia', 'cyclamen ', 'watercress',  'canna lily', \n    'hippeastrum ', 'bee balm', 'pink quill',  'foxglove', 'bougainvillea', \n    'camellia', 'mallow',  'mexican petunia',  'bromelia', 'blanket flower', \n    'trumpet creeper',  'blackberry lily', 'common tulip', 'wild rose']","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Helper and Augmentation functions"},{"metadata":{"_kg_hide-input":false,"trusted":true},"cell_type":"code","source":"# Datasets utility functions\nAUTO = tf.data.experimental.AUTOTUNE # instructs the API to read from multiple files if available.\n\ndef decode_image(image_data):\n    image = tf.image.decode_jpeg(image_data, channels=3)\n    image = tf.cast(image, tf.float32) / 255.0\n    image = tf.reshape(image, [HEIGHT, WIDTH, 3])\n    return image\n\ndef read_labeled_tfrecord(example):\n    LABELED_TFREC_FORMAT = {\n        \"image\": tf.io.FixedLenFeature([], tf.string), # tf.string means bytestring\n        \"class\": tf.io.FixedLenFeature([], tf.int64),  # shape [] means single element\n    }\n    example = tf.io.parse_single_example(example, LABELED_TFREC_FORMAT)\n    image = decode_image(example['image'])\n    label = tf.cast(example['class'], tf.int32)\n    return image, label\n\ndef read_unlabeled_tfrecord(example):\n    UNLABELED_TFREC_FORMAT = {\n        \"image\": tf.io.FixedLenFeature([], tf.string), # tf.string means bytestring\n        \"id\": tf.io.FixedLenFeature([], tf.string),  # shape [] means single element\n        # class is missing, this competitions's challenge is to predict flower classes for the test dataset\n    }\n    example = tf.io.parse_single_example(example, UNLABELED_TFREC_FORMAT)\n    image = decode_image(example['image'])\n    idnum = example['id']\n    return image, idnum # returns a dataset of image(s)\n\ndef load_dataset(filenames, labeled=True, ordered=False):\n    ignore_order = tf.data.Options()\n    if not ordered:\n        ignore_order.experimental_deterministic = False # disable order, increase speed\n\n    dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=AUTO) # automatically interleaves reads from multiple files\n    dataset = dataset.with_options(ignore_order) # uses data as soon as it streams in, rather than in its original order\n    dataset = dataset.map(read_labeled_tfrecord if labeled else read_unlabeled_tfrecord, num_parallel_calls=AUTO)\n    # returns a dataset of (image, label) pairs if labeled=True or (image, id) pairs if labeled=False\n    return dataset\n\ndef data_augment(image, label):\n    \n    # Take a random crop of 480x480 and then apply random horizontal flip\n    image = tf.image.random_crop(image, size=[480, 480, CHANNELS])\n    image = tf.image.random_flip_left_right(image, seed=SEED)\n\n    return image, label\n\ndef get_training_dataset():\n    dataset = load_dataset(TRAINING_FILENAMES, labeled=True)\n    dataset = dataset.map(data_augment, num_parallel_calls=AUTO)\n    dataset = dataset.repeat() # the training dataset must repeat for several epochs\n    dataset = dataset.shuffle(2048)\n    dataset = dataset.batch(BATCH_SIZE)\n    dataset = dataset.prefetch(AUTO) # prefetch next batch while training (autotune prefetch buffer size)\n    return dataset\n\ndef get_validation_dataset(ordered=False):\n    dataset = load_dataset(VALIDATION_FILENAMES, labeled=True, ordered=ordered)\n    dataset = dataset.batch(BATCH_SIZE)\n    dataset = dataset.cache()\n    dataset = dataset.prefetch(AUTO)\n    return dataset\n\ndef get_train_valid_datasets():\n    dataset = load_dataset(TRAINING_FILENAMES + VALIDATION_FILENAMES, labeled=True)\n    dataset = dataset.map(data_augment, num_parallel_calls=AUTO)\n    dataset = dataset.repeat() # the training dataset must repeat for several epochs\n    dataset = dataset.shuffle(2048)\n    dataset = dataset.batch(BATCH_SIZE)\n    dataset = dataset.prefetch(AUTO) # prefetch next batch while training (autotune prefetch buffer size)\n    return dataset\n\ndef get_test_dataset(ordered=False):\n    dataset = load_dataset(TEST_FILENAMES, labeled=False, ordered=ordered)\n    dataset = dataset.batch(BATCH_SIZE)\n    dataset = dataset.prefetch(AUTO)\n    return dataset\n\ndef count_data_items(filenames):\n    # the number of data items is written in the name of the .tfrec files, i.e. flowers00-230.tfrec = 230 data items\n    n = [int(re.compile(r\"-([0-9]*)\\.\").search(filename).group(1)) for filename in filenames]\n    return np.sum(n)","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"# Visualization utility functions\nnp.set_printoptions(threshold=15, linewidth=80)\n\ndef batch_to_numpy_images_and_labels(data):\n    images, labels = data\n    numpy_images = images.numpy()\n    numpy_labels = labels.numpy()\n    if numpy_labels.dtype == object: # binary string in this case, these are image ID strings\n        numpy_labels = [None for _ in enumerate(numpy_images)]\n    # If no labels, only image IDs, return None for labels (this is the case for test data)\n    return numpy_images, numpy_labels\n\ndef title_from_label_and_target(label, correct_label):\n    if correct_label is None:\n        return CLASSES[label], True\n    correct = (label == correct_label)\n    return \"{} [{}{}{}]\".format(CLASSES[label], 'OK' if correct else 'NO', u\"\\u2192\" if not correct else '',\n                                CLASSES[correct_label] if not correct else ''), correct\n\ndef display_one_flower(image, title, subplot, red=False, titlesize=16):\n    plt.subplot(*subplot)\n    plt.axis('off')\n    plt.imshow(image)\n    if len(title) > 0:\n        plt.title(title, fontsize=int(titlesize) if not red else int(titlesize/1.2), color='red' if red else 'black', fontdict={'verticalalignment':'center'}, pad=int(titlesize/1.5))\n    return (subplot[0], subplot[1], subplot[2]+1)\n\ndef display_batch_of_images(databatch, predictions=None):\n    \"\"\"This will work with:\n    display_batch_of_images(images)\n    display_batch_of_images(images, predictions)\n    display_batch_of_images((images, labels))\n    display_batch_of_images((images, labels), predictions)\n    \"\"\"\n    # data\n    images, labels = batch_to_numpy_images_and_labels(databatch)\n    if labels is None:\n        labels = [None for _ in enumerate(images)]\n        \n    # auto-squaring: this will drop data that does not fit into square or square-ish rectangle\n    rows = int(math.sqrt(len(images)))\n    cols = len(images)//rows\n        \n    # size and spacing\n    FIGSIZE = 13.0\n    SPACING = 0.1\n    subplot=(rows,cols,1)\n    if rows < cols:\n        plt.figure(figsize=(FIGSIZE,FIGSIZE/cols*rows))\n    else:\n        plt.figure(figsize=(FIGSIZE/rows*cols,FIGSIZE))\n    \n    # display\n    for i, (image, label) in enumerate(zip(images[:rows*cols], labels[:rows*cols])):\n        title = '' if label is None else CLASSES[label]\n        correct = True\n        if predictions is not None:\n            title, correct = title_from_label_and_target(predictions[i], label)\n        dynamic_titlesize = FIGSIZE*SPACING/max(rows,cols)*40+3 # magic formula tested to work from 1x1 to 10x10 images\n        subplot = display_one_flower(image, title, subplot, not correct, titlesize=dynamic_titlesize)\n    \n    #layout\n    plt.tight_layout()\n    if label is None and predictions is None:\n        plt.subplots_adjust(wspace=0, hspace=0)\n    else:\n        plt.subplots_adjust(wspace=SPACING, hspace=SPACING)\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"def display_training_curves(training, validation, title, subplot):\n    if subplot%10==1: # set up the subplots on the first call\n        plt.subplots(figsize=(15,10), facecolor='#F0F0F0')\n        plt.tight_layout()\n    ax = plt.subplot(subplot)\n    ax.set_facecolor('#F8F8F8')\n    ax.plot(training)\n    ax.plot(validation)\n    ax.set_title(title)\n    ax.set_ylabel(title.split(\" \")[-1])\n    #ax.set_ylim(0.28,1.05)\n    ax.set_xlabel('epoch')\n    ax.legend(['train', 'valid.'])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Count train, val and test data\nNUM_TRAINING_IMAGES = count_data_items(TRAINING_FILENAMES)\nNUM_VALIDATION_IMAGES = count_data_items(VALIDATION_FILENAMES)\nNUM_TEST_IMAGES = count_data_items(TEST_FILENAMES)\n\nprint('Number of training images %d' % NUM_TRAINING_IMAGES)\nprint('Number of validation images %d' % NUM_VALIDATION_IMAGES)\nprint('Number of test images %d' % NUM_TEST_IMAGES)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Augmentation Visualization\n- random crop of 480x480 \n- random horizontal flip"},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset = load_dataset(TRAINING_FILENAMES, labeled=True)\none_element = tf.data.Dataset.from_tensors( next(iter(dataset)) )\naugmented_element = one_element.repeat().map(data_augment).batch(25)\n\ndisplay_batch_of_images(next(iter(augmented_element)))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Model function and training configurations"},{"metadata":{"trusted":true},"cell_type":"code","source":"class BiTModel(tf.keras.Model):\n\n    def __init__(self, num_classes, module):\n        super().__init__()\n\n        self.num_classes = num_classes\n        self.head = tf.keras.layers.Dense(num_classes, kernel_initializer='zeros')\n        self.bit_model = module\n\n    def call(self, images):\n        bit_embedding = self.bit_model(images)\n        return self.head(bit_embedding)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"LR = 3e-3\nSCHEDULE_LENGTH = 500\nSTEPS_PER_EPOCH = 50\nBATCH_SIZE = 32 * strategy.num_replicas_in_sync\n\n\nLR = 0.003 * BATCH_SIZE / 512 \nSCHEDULE_LENGTH = SCHEDULE_LENGTH * 512 / BATCH_SIZE\nEPOCHS = int(SCHEDULE_LENGTH / STEPS_PER_EPOCH)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"with strategy.scope():\n    GCS_PATH_TO_SAVEDMODEL = KaggleDatasets().get_gcs_path('bit-tf-hub-models') + '/BiT TF Hub Models/bit_m-r101x1_1'\n    loaded_model = tf.saved_model.load(GCS_PATH_TO_SAVEDMODEL)\n    model = BiTModel(num_classes=N_CLASSES, module=loaded_model)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"SCHEDULE_BOUNDARIES = [int(SCHEDULE_LENGTH * 0.3), int(SCHEDULE_LENGTH * 0.6), int(SCHEDULE_LENGTH * 0.9)]\n\n# Decay learning rate by a factor of 10 at SCHEDULE_BOUNDARIES.\nlr_schedule = tf.keras.optimizers.schedules \\\n                .PiecewiseConstantDecay(\n                    boundaries=SCHEDULE_BOUNDARIES, values=[LR, LR*0.1, LR*0.001, LR*0.0001]\n                )\n\noptimizer = tf.keras.optimizers.SGD(learning_rate=lr_schedule, momentum=0.9)\nloss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n\nmodel.compile(\n    optimizer=optimizer,\n    loss = loss_fn,\n    metrics=['sparse_categorical_accuracy']\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"history = model.fit(x=get_training_dataset(),\n                    steps_per_epoch=STEPS_PER_EPOCH,\n                    validation_data=get_validation_dataset(),\n                    epochs=EPOCHS, verbose=1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"val_ds = get_validation_dataset(ordered=True) # since we are splitting the dataset and iterating separately on images and labels, order matters.\nval_images_ds = val_ds.map(lambda image, label: image)\nval_labels_ds = val_ds.map(lambda image, label: label).unbatch()\nval_correct_labels = next(iter(val_labels_ds.batch(NUM_VALIDATION_IMAGES))).numpy() # get everything as one batch\n\ntest_ds = get_test_dataset(ordered=True) # since we are splitting the dataset and iterating separately on images and ids, order matters.\ntest_images_ds = test_ds.map(lambda image, idnum: image)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"val_probabilities = model.predict(val_images_ds)\nval_predictions = np.argmax(val_probabilities, axis=-1)\n\nscore_sch = f1_score(val_correct_labels, val_predictions, labels=range(len(CLASSES)), average='macro')\nprecision_sch = precision_score(val_correct_labels, val_predictions, labels=range(len(CLASSES)), average='macro')\nrecall_sch = recall_score(val_correct_labels, val_predictions, labels=range(len(CLASSES)), average='macro')\n\ntest_probabilities = model.predict(test_images_ds)\n\nprint('f1 score: {:.3f}, precision: {:.3f}, recall: {:.3f}'.format(score_sch, precision_sch, recall_sch))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"display_training_curves(history.history['loss'], history.history['val_loss'], 'BiT Loss', 221)\ndisplay_training_curves(history.history['sparse_categorical_accuracy'], history.history['val_sparse_categorical_accuracy'], 'BiT accuracy', 222)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print('[INFO] Computing predictions...')\n\nprobabilities = test_probabilities\npredictions = np.argmax(probabilities, axis=-1)\n\nprint('[INFO] Generating submission.csv file...')\ntest_ids_ds = test_ds.map(lambda image, idnum: idnum).unbatch()\ntest_ids = next(iter(test_ids_ds.batch(NUM_TEST_IMAGES))).numpy().astype('U') # all in one batch\nnp.savetxt('submission.csv', np.rec.fromarrays([test_ids, predictions]), fmt=['%s', '%d'], delimiter=',', header='id,label', comments='')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","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}