{"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":"## 1. Import required libraries","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19"}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport tensorflow as tf\nimport tensorflow.keras.layers as L\nimport tensorflow_addons as tfa\nimport glob, random, os, warnings\nimport matplotlib.pyplot as plt\nfrom sklearn.metrics import confusion_matrix, classification_report\nimport seaborn as sns\n\nprint('TensorFlow Version ' + tf.__version__)\n\n# setting random seed\nseed = 0\nrandom.seed(seed)\nnp.random.seed(seed)\ntf.random.set_seed(seed)\nos.environ['PYTHONHASHSEED'] = str(seed)\nos.environ['TF_DETERMINISTIC_OPS'] = '1'","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 2. Setting input file paths and labels","metadata":{}},{"cell_type":"code","source":"image_size = 224\nbatch_size = 16\nn_classes = 5\n\ntrain_path = '/kaggle/input/cassava-leaf-disease-classification/train_images'\ntest_path = '/kaggle/input/cassava-leaf-disease-classification/test_images'\n\ndf_train = pd.read_csv('/kaggle/input/cassava-leaf-disease-classification/train.csv', dtype = 'str')\n\ntest_images = glob.glob(test_path + '/*.jpg')\ndf_test = pd.DataFrame(test_images, columns = ['image_path'])\n\nclasses = {0 : \"Cassava Bacterial Blight (CBB)\",\n           1 : \"Cassava Brown Streak Disease (CBSD)\",\n           2 : \"Cassava Green Mottle (CGM)\",\n           3 : \"Cassava Mosaic Disease (CMD)\",\n           4 : \"Healthy\"}","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"","metadata":{}},{"cell_type":"markdown","source":"## 3. Data preparation","metadata":{}},{"cell_type":"code","source":"datagen = tf.keras.preprocessing.image.ImageDataGenerator(samplewise_center = True,\n                                                          samplewise_std_normalization = True,\n                                                          validation_split = 0.2)\n\ntrain_gen = datagen.flow_from_dataframe(dataframe = df_train,\n                                        directory = train_path,\n                                        x_col = 'image_id',\n                                        y_col = 'label',\n                                        subset = 'training',\n                                        batch_size = batch_size,\n                                        seed = 1,\n                                        color_mode = 'rgb',\n                                        shuffle = True,\n                                        class_mode = 'categorical',\n                                        target_size = (image_size, image_size))\n\nvalid_gen = datagen.flow_from_dataframe(dataframe = df_train,\n                                        directory = train_path,\n                                        x_col = 'image_id',\n                                        y_col = 'label',\n                                        subset = 'validation',\n                                        batch_size = batch_size,\n                                        seed = 1,\n                                        color_mode = 'rgb',\n                                        shuffle = False,\n                                        class_mode = 'categorical',\n                                        target_size = (image_size, image_size))\n\ntest_gen = datagen.flow_from_dataframe(dataframe = df_test,\n                                       x_col = 'image_path',\n                                       y_col = None,\n                                       batch_size = batch_size,\n                                       seed = 1,\n                                       color_mode = 'rgb',\n                                       shuffle = False,\n                                       class_mode = None,\n                                       target_size = (image_size, image_size))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 4. Visualize sample data","metadata":{}},{"cell_type":"code","source":"images = [train_gen[0][0][i] for i in range(10, 14)]\nfig, axes = plt.subplots(1, 4, figsize = (10, 10))\n\naxes = axes.flatten()\n\nfor img, ax in zip(images, axes):\n    ax.imshow(img.reshape(image_size, image_size, 3))\n    ax.axis('off')\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 5. Setting hyperparameters","metadata":{}},{"cell_type":"code","source":"learning_rate = 0.001\nweight_decay = 0.0001\nnum_epochs = 1\n\npatch_size = 7  # Size of the patches to be extract from the input images\nnum_patches = (image_size // patch_size) ** 2\nprojection_dim = 64\nnum_heads = 4\ntransformer_units = [\n    projection_dim * 2,\n    projection_dim,\n]  # Size of the transformer layers\ntransformer_layers = 8\nmlp_head_units = [56, 28]  # Size of the dense layers of the final classifier","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 6. Defining ViT model and its components","metadata":{}},{"cell_type":"markdown","source":"### 6.1 Define MLP layer","metadata":{}},{"cell_type":"code","source":"# MLP layer\ndef mlp(x, hidden_units, dropout_rate):\n    for units in hidden_units:\n        x = L.Dense(units, activation = tf.nn.gelu)(x)\n        x = L.Dropout(dropout_rate)(x)\n    return x","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 6.2 Define Patch creation layer\n- The image must be fed as patches to the encoder","metadata":{}},{"cell_type":"code","source":"# Patch creation layer\nclass Patches(L.Layer):\n    def __init__(self, patch_size):\n        super(Patches, self).__init__()\n        self.patch_size = patch_size\n\n    def call(self, images):\n        batch_size = tf.shape(images)[0]\n        patches = tf.image.extract_patches(\n            images = images,\n            sizes = [1, self.patch_size, self.patch_size, 1],\n            strides = [1, self.patch_size, self.patch_size, 1],\n            rates = [1, 1, 1, 1],\n            padding = 'VALID',\n        )\n        patch_dims = patches.shape[-1]\n        patches = tf.reshape(patches, [batch_size, -1, patch_dims])\n        return patches","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"- Visualize how an image in divided into patches before feeding to encoder network","metadata":{}},{"cell_type":"code","source":"plt.figure(figsize=(4, 4))\n\nx = train_gen[10]\nimage = x[0][0]\n\nplt.imshow(image)\nplt.axis('off')\n\nresized_image = tf.image.resize(\n    tf.convert_to_tensor([image]), size = (image_size, image_size)\n)\n\npatches = Patches(patch_size)(resized_image)\nprint(f'Image size: {image_size} X {image_size}')\nprint(f'Patch size: {patch_size} X {patch_size}')\nprint(f'Patches per image: {patches.shape[1]}')\nprint(f'Elements per patch: {patches.shape[-1]}')\n\nn = int(np.sqrt(patches.shape[1]))\nplt.figure(figsize=(4, 4))\n\nfor i, patch in enumerate(patches[0]):\n    ax = plt.subplot(n, n, i + 1)\n    patch_img = tf.reshape(patch, (patch_size, patch_size, 3))\n    plt.imshow(patch_img.numpy())\n    plt.axis('off')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 6.3 Define Patch encoding layer\n- The image patches are unrolled into a vector\n- Then image vector embeddings are generated\n- Positional encoding is also added to these image patch embeddings","metadata":{}},{"cell_type":"code","source":"class PatchEncoder(L.Layer):\n    def __init__(self, num_patches, projection_dim):\n        super(PatchEncoder, self).__init__()\n        self.num_patches = num_patches\n        self.projection = L.Dense(units = projection_dim)\n        self.position_embedding = L.Embedding(\n            input_dim = num_patches, output_dim = projection_dim\n        )\n\n    def call(self, patch):\n        positions = tf.range(start = 0, limit = self.num_patches, delta = 1)\n        encoded = self.projection(patch) + self.position_embedding(positions)\n        return encoded","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### 6.4 Define ViT model","metadata":{}},{"cell_type":"code","source":"def vision_transformer():\n    inputs = L.Input(shape = (image_size, image_size, 3))\n    \n    # Create patches.\n    patches = Patches(patch_size)(inputs)\n    \n    # Encode patches.\n    encoded_patches = PatchEncoder(num_patches, projection_dim)(patches)\n\n    # Create multiple layers of the Transformer block.\n    for _ in range(transformer_layers):\n        \n        # Layer normalization 1.\n        x1 = L.LayerNormalization(epsilon = 1e-6)(encoded_patches)\n        \n        # Create a multi-head attention layer.\n        attention_output = L.MultiHeadAttention(\n            num_heads = num_heads, key_dim = projection_dim, dropout = 0.1\n        )(x1, x1)\n        \n        # Skip connection 1.\n        x2 = L.Add()([attention_output, encoded_patches])\n        \n        # Layer normalization 2.\n        x3 = L.LayerNormalization(epsilon = 1e-6)(x2)\n        \n        # MLP.\n        x3 = mlp(x3, hidden_units = transformer_units, dropout_rate = 0.1)\n        \n        # Skip connection 2.\n        encoded_patches = L.Add()([x3, x2])\n\n    # Create a [batch_size, projection_dim] tensor.\n    representation = L.LayerNormalization(epsilon = 1e-6)(encoded_patches)\n    representation = L.Flatten()(representation)\n    representation = L.Dropout(0.5)(representation)\n    \n    # Add MLP.\n    features = mlp(representation, hidden_units = mlp_head_units, dropout_rate = 0.5)\n    \n    # Classify outputs.\n    logits = L.Dense(n_classes)(features)\n    \n    # Create the model.\n    model = tf.keras.Model(inputs = inputs, outputs = logits)\n    \n    return model","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 7. Train the ViT model","metadata":{}},{"cell_type":"code","source":"decay_steps = train_gen.n // train_gen.batch_size\ninitial_learning_rate = learning_rate\n\nlr_decayed_fn = tf.keras.experimental.CosineDecay(initial_learning_rate, decay_steps)\n\nlr_scheduler = tf.keras.callbacks.LearningRateScheduler(lr_decayed_fn)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"gpus = tf.config.experimental.list_physical_devices('GPU')\nif gpus:\n  try:\n    tf.config.experimental.set_virtual_device_configuration(\n        gpus[0],[tf.config.experimental.VirtualDeviceConfiguration(memory_limit=5120)])\n  except RuntimeError as e:\n    print(e)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"optimizer = tf.keras.optimizers.Adam(learning_rate = learning_rate)\n\nmodel = vision_transformer()\n    \nmodel.compile(optimizer = optimizer, \n              loss = tf.keras.losses.CategoricalCrossentropy(label_smoothing = 0.1), \n              metrics = ['accuracy'])\n\n\nSTEP_SIZE_TRAIN = train_gen.n // train_gen.batch_size\nSTEP_SIZE_VALID = valid_gen.n // valid_gen.batch_size\n\nearlystopping = tf.keras.callbacks.EarlyStopping(monitor = 'val_accuracy',\n                                                 min_delta = 1e-4,\n                                                 patience = 5,\n                                                 mode = 'max',\n                                                 restore_best_weights = True,\n                                                 verbose = 1)\n\ncheckpointer = tf.keras.callbacks.ModelCheckpoint(filepath = './model.hdf5',\n                                                  monitor = 'val_accuracy', \n                                                  verbose = 1, \n                                                  save_best_only = True,\n                                                  save_weights_only = True,\n                                                  mode = 'max')\n\ncallbacks = [earlystopping, lr_scheduler, checkpointer]\n\nmodel.fit(x = train_gen,\n          steps_per_epoch = STEP_SIZE_TRAIN,\n          validation_data = valid_gen,\n          validation_steps = STEP_SIZE_VALID,\n          epochs = num_epochs,\n          callbacks = callbacks)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 8. Evaluate the trained ViT model","metadata":{}},{"cell_type":"code","source":"print('Training results')\nmodel.evaluate(train_gen)\n\nprint('Validation results')\nmodel.evaluate(valid_gen)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### Reference\n- https://www.kaggle.com/raufmomin/vision-transformer-vit-from-scratch/notebook","metadata":{}}]}