{"cells":[{"metadata":{},"cell_type":"markdown","source":"## Intro\n**Hello kagglers.\nIn this notebook I will load, train and save TFHub models with TPU. It's not for LB points, just for study.**\n**Some code was taken from this** [notebook](https://www.kaggle.com/smirnyaginandr/notebooke180f47a7c)"},{"metadata":{},"cell_type":"markdown","source":"## Step Uno\n**Just import**"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import pandas as pd \nimport os\nfrom functools import partial\nimport tensorflow as tf\nimport re, math\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport pandas as pd\nfrom kaggle_datasets import KaggleDatasets\nimport csv\nimport tensorflow_hub as hub\nfrom sklearn.model_selection import train_test_split\nfrom kaggle_secrets import UserSecretsClient\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Init datasets for use with TPU. Dataset 'tf-hub' contains 4 models from TFHub. I will use them for training.**"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"GCS_DS_PATH = KaggleDatasets().get_gcs_path('cassava-leaf-disease-classification')\nGCS_PATH_TO_SAVEDMODEL = KaggleDatasets().get_gcs_path('tf-hub')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Step zwei\n**Init resolver and strategy. It MUST be before the next step!**"},{"metadata":{"trusted":true},"cell_type":"code","source":"try:\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() \nprint(\"REPLICAS: \", strategy.num_replicas_in_sync)\nAUTO = tf.data.experimental.AUTOTUNE","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Step three\n**Set up credentials to grant kaggle access to your Google Cloud account. It's important for reading private datasets and for saving trained models to cloud storage.\nThis step MUST be after the init TPU strategy. Thanks to @morodertobias for the advice.**"},{"metadata":{"trusted":true},"cell_type":"code","source":"user_secrets = UserSecretsClient()\nuser_credential = user_secrets.get_gcloud_credential()\nuser_secrets.set_tensorflow_credential(user_credential)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Step четыре\n**Get filenames for train and validation sets**"},{"metadata":{"trusted":true},"cell_type":"code","source":"TRAINING_FILENAMES, VALID_FILENAMES = train_test_split(\n    tf.io.gfile.glob(GCS_DS_PATH + '/train_tfrecords/ld_train*.tfrec'),\n    test_size=0.35, random_state=5\n)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Define some variables**"},{"metadata":{"trusted":true},"cell_type":"code","source":"BATCH_SIZE = 16 * strategy.num_replicas_in_sync\nAUTOTUNE = tf.data.experimental.AUTOTUNE\nCLASSES = 5\nCHANNELS = 3\nSEED = 42\nDIM = 224","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Define functions to augment data and create datasets**"},{"metadata":{"trusted":true},"cell_type":"code","source":"def 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.image.resize(image, [DIM, DIM])\n    return image\n\n\ndef read_labeled_tfrecord(example):\n    LABELED_TFREC_FORMAT = {\n        \"image\"              : tf.io.FixedLenFeature([], tf.string),\n        \"target\"              : tf.io.FixedLenFeature([], tf.int64),\n    }\n    \n    example = tf.io.parse_single_example(example, LABELED_TFREC_FORMAT)\n    image = decode_image(example['image'])\n    label = example[\"target\"]\n    label = tf.one_hot(label,depth=5)\n    label = tf.cast(label,tf.float32)\n    return image, label\n\n\ndef read_unlabeled_tfrecord(example):\n    UNLABELED_TFREC_FORMAT = {\n        \"image\"              : tf.io.FixedLenFeature([], tf.string),\n        \"image_name\"           : tf.io.FixedLenFeature([], tf.string),\n    }\n    example = tf.io.parse_single_example(example, UNLABELED_TFREC_FORMAT)\n    image = decode_image(example['image'])\n    idnum = example['image_name']\n    return image, idnum\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 \n    dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=AUTO) \n    dataset = dataset.with_options(ignore_order)\n    dataset = dataset.map(read_labeled_tfrecord if labeled else read_unlabeled_tfrecord, num_parallel_calls=AUTO)\n    return dataset\n\ndef data_augment(image, label):\n    p_spatial = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    p_rotate = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    p_pixel_1 = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    p_pixel_2 = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    p_pixel_3 = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n    p_crop = tf.random.uniform([], 0, 1.0, dtype=tf.float32)\n            \n    # Flips\n    image = tf.image.random_flip_left_right(image)\n    image = tf.image.random_flip_up_down(image)\n    if p_spatial > .75:\n        image = tf.image.transpose(image)\n        \n    # Rotates\n    if p_rotate > .75:\n        image = tf.image.rot90(image, k=3) # rotate 270º\n    elif p_rotate > .5:\n        image = tf.image.rot90(image, k=2) # rotate 180º\n    elif p_rotate > .25:\n        image = tf.image.rot90(image, k=1) # rotate 90º\n        \n    # Pixel-level transforms\n    if p_pixel_1 >= .4:\n        image = tf.image.random_saturation(image, lower=.7, upper=1.3)\n    if p_pixel_2 >= .4:\n        image = tf.image.random_contrast(image, lower=.8, upper=1.2)\n    if p_pixel_3 >= .4:\n        image = tf.image.random_brightness(image, max_delta=.1)\n\n    return image, label  \n\ndef get_training_dataset():\n    dataset = load_dataset(TRAINING_FILENAMES, labeled=True)\n    dataset = dataset.map(data_augment)\n    dataset = dataset.repeat() \n    dataset = dataset.shuffle(2048)\n    dataset = dataset.batch(BATCH_SIZE)\n    dataset = dataset.prefetch(AUTOTUNE)\n    return dataset\n\ndef get_validation_dataset(ordered=False):\n    dataset = load_dataset(VALID_FILENAMES, labeled=True, ordered=ordered)\n    dataset = dataset.batch(BATCH_SIZE)\n    dataset = dataset.cache()\n    dataset = dataset.prefetch(AUTOTUNE)\n    return dataset","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Get datasets**"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dataset = get_training_dataset()\nvalidation_dataset = get_validation_dataset()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Count steps for train and validation**"},{"metadata":{"trusted":true},"cell_type":"code","source":"def count_data_items(filenames):\n    n = [int(re.compile(r\"-([0-9]*)\\.\").search(filename).group(1)) for filename in filenames]\n    return np.sum(n)\n\nSTEPS_PER_EPOCH = count_data_items(TRAINING_FILENAMES) // BATCH_SIZE\nVALID_STEPS = count_data_items(VALID_FILENAMES) // BATCH_SIZE","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**Define callbacks**"},{"metadata":{"trusted":true},"cell_type":"code","source":"early_stop = tf.keras.callbacks.EarlyStopping(monitor = 'val_acc', min_delta = 0.001, \n                           patience = 10, mode = 'max', verbose = 1,\n                           restore_best_weights = True)\nreduce_lr = tf.keras.callbacks.ReduceLROnPlateau(monitor = 'val_loss', factor = 0.1, \n                              patience = 2, min_delta = 0.001, \n                              mode = 'min', verbose = 1)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Step 五\n**Before this step you need creata a bucket on the Google Cloud.** \n\n![](https://storage.cloud.google.com/cassava_saved_models/cr_bucket-1.png)"},{"metadata":{},"cell_type":"markdown","source":"**I called my bucket \"cassava_saved_models\".**\n\n![](https://storage.cloud.google.com/cassava_saved_models/cr_bucket-2.png)"},{"metadata":{},"cell_type":"markdown","source":"**And then just train the models. I am skipping \"efficientnet\" because this model contains a module for TF1 and unfortunately I havn't found a way to load this model correctly.**"},{"metadata":{"trusted":true},"cell_type":"code","source":"tf.keras.backend.clear_session()\nmodels = [GCS_PATH_TO_SAVEDMODEL + '/nets/' + x for x in tf.io.gfile.listdir(GCS_PATH_TO_SAVEDMODEL + '/nets/')]\nwith strategy.scope():\n    \n    for model_path in models:\n        model_name = model_path.split('/')[-2]\n        if model_name.startswith('efficientnet'):\n            continue\n        print()\n        print('=' * 100)\n        print(f'Load model {model_name}')\n        \n        loaded_model = tf.saved_model.load(model_path)\n        base_model = hub.KerasLayer(loaded_model, trainable=True)\n        model = tf.keras.Sequential([\n            tf.keras.Input(shape=(DIM, DIM, 3)),\n            base_model, \n            tf.keras.layers.Dense(5, activation='softmax')\n        ])\n        model.compile(\n                optimizer=tf.keras.optimizers.Adam(),\n                loss='categorical_crossentropy',  \n                metrics=['acc'])\n\n        history = model.fit(train_dataset, \n                            epochs=20,\n                            callbacks=[early_stop, reduce_lr],\n                            validation_data = validation_dataset,\n                            steps_per_epoch = STEPS_PER_EPOCH,\n                            verbose=1)\n        model.save(f'gs://cassava_saved_models/saved_{model_name}_{int(max(history.history[\"val_acc\"])*100)}')\n        print(f'Model {model_name} saved')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**After training you will find saved model in your bucket.**\n![](https://storage.cloud.google.com/cassava_saved_models/cr_bucket-3.png)"}],"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}