{"cells":[{"metadata":{},"cell_type":"markdown","source":"![](https://hubmapconsortium.github.io/ccf/img/img-spatial-steps-registration-01.png)\n\n

Figure: Spatial placement of anatomical structures in relation to the HuBMAP Atlas reference system.

\n

\n\nRESUME TPU KERAS TRAINING FOR LARGE MODELS\n\n

This notebook contains steps to resume TPU training from previous saved model weights. This is useful for large models which needs to be trained more than the cap limit per session in Kaggle Notebooks

\n\n\n*

The model weights and the callback files are automatically uploaded to GCS at the end of every epoch by means of a Custom Callback Model Checkpoint Function.

\n \n \n\n*

On the first run, a GCS bucket is created and session files are stored here. On any subsequent runs, the latest model weights and callback files in the GCS bucket are used to resume training.

\n\n\n\n*

Once the session expires, simply re-running the whole notebook again will resume the training from the last saved epoch.

\n\n\nUseful Datasets:\n\n\nHubMAP TFRecords With Augmentation Dataset is being used in this notebook. It has TFRecords with actual and augmented images grouped into 90-130MB records fit for TPU loads.\n\n - HubMAP TFRecords With Augmentation Dataset: https://www.kaggle.com/sreevishnudamodaran/hubmap-512x512-tfrecords-with-aug\n\n\nCreated from **Augmented images 512x512 tiled Dataset** which has augmented images and masks.\n\n - HuBMAP Augmented 512x512: https://www.kaggle.com/sreevishnudamodaran/hubmap-512x512-augmented\n\nPlease take a look at my previous notebook on creating these datasets and building a Double U-net model:\n\n - TPU - HubMAP Double U-Net Model + Augmentation : https://www.kaggle.com/sreevishnudamodaran/tpu-hubmap-double-u-net-model-augmentation\n\n

I would like to highlight that we should make use of TPUs judiciously and make sure that it is available for everyone. Thanks to the Kaggle team for such great perks!

\n\nNote:\n*

Please note that it may not be possible to recreate the exact conditions of previous training session as it is not possible to keep track of the state of certain callbacks at present. Pickling of such callbacks are not supported at the moment due to thread.R_lock, thread.lock and lambda fn objects in the tensorflow code.

\n \n \n*

For resuming training, a GCS Storage Bucket will be created/used as a part of this notebook. So a GCP account and a project has to be setup prior to running the notebook. The Google Cloud SDK and Google Cloud Services also has to be enabled from the Add-ons menu.

\n\n \n \n\n[![Ask Me Anything !](https://img.shields.io/badge/Ask%20me-anything-1abc9c.svg?style=flat-square&logo=appveyor)](https://www.kaggle.com/sreevishnudamodaran)\n\n\n\n![TPU!](https://img.shields.io/badge/Accelerator-TPU-green?style=flat-square&logo=appveyor)\n\n![Upvote!](https://img.shields.io/badge/Upvote-If%20you%20like%20my%20work-blue?style=for-the-badge&logo=appveyor)"},{"metadata":{},"cell_type":"markdown","source":"## Import Libraries"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"%matplotlib inline\n\nimport json\nimport os\nimport glob\nimport re\nimport datetime\nimport os.path as osp\nfrom path import Path\nimport collections\nimport sys\nimport uuid\nimport random\nimport warnings\nfrom itertools import chain\nfrom tqdm import tqdm\n\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nfrom matplotlib import rcParams\nsns.set(rc={\"font.size\":9,\"axes.titlesize\":15,\"axes.labelsize\":9,\n \"axes.titlepad\":2, \"axes.labelpad\":9, \"legend.fontsize\":7,\n \"legend.title_fontsize\":7, 'axes.grid' : False,\n 'figure.titlesize':35\n \n })\n\nfrom PIL import Image\nimport cv2\n\nimport tensorflow as tf\nfrom tensorflow.keras.models import Model, load_model\nfrom tensorflow.keras.layers import Input\nfrom tensorflow.keras.layers import Conv2D, Conv2DTranspose\nfrom tensorflow.keras.layers import MaxPooling2D, UpSampling2D, Conv2DTranspose\nfrom tensorflow.keras.layers import concatenate\nfrom tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint\nfrom tensorflow.keras import backend as K\nfrom tensorflow.keras import layers\nfrom tensorflow.keras.optimizers import Adam\nfrom tensorflow.keras.utils import get_custom_objects\nfrom tensorflow.keras.losses import binary_crossentropy\n\nfrom kaggle_datasets import KaggleDatasets\nfrom kaggle_secrets import UserSecretsClient","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Intialize and Get TPU Ready"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\ntf.config.experimental_connect_to_cluster(tpu)\ntf.tpu.experimental.initialize_tpu_system(tpu)\n\nstrategy = tf.distribute.experimental.TPUStrategy(tpu)\n","execution_count":null,"outputs":[]},{"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":"## Load TFRecords & View Samples"},{"metadata":{"trusted":true},"cell_type":"code","source":"AUTO = tf.data.experimental.AUTOTUNE\n\nimage_feature_description = {\n 'height': tf.io.FixedLenFeature([], tf.int64),\n 'width': tf.io.FixedLenFeature([], tf.int64),\n 'num_channels': tf.io.FixedLenFeature([], tf.int64),\n 'img_bytes': tf.io.FixedLenFeature([], tf.string),\n 'mask': tf.io.FixedLenFeature([], tf.string),\n}\n\ndef _parse_image_and_masks_function(example_proto):\n single_example = tf.io.parse_single_example(example_proto, image_feature_description)\n img_bytes = tf.io.decode_raw(single_example['img_bytes'], out_type='uint8')\n img_array = tf.reshape(img_bytes, (512, 512, 3))\n mask_bytes = tf.io.decode_raw(single_example['mask'], out_type='bool')\n mask = tf.reshape(mask_bytes, (512, 512, 1))\n \n ## normalize images array and cast image and mask to float32\n# img_array = tf.cast(img_array, tf.float32) / 255.0\n# mask = tf.cast(mask, tf.float32)\n return img_array, mask\n\ndef read_dataset(storage_file_path):\n encoded_image_dataset = tf.data.TFRecordDataset(storage_file_path, compression_type=\"GZIP\")\n parsed_image_dataset = encoded_image_dataset.map(_parse_image_and_masks_function)\n return parsed_image_dataset","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Get the credential from the Cloud SDK\nfrom kaggle_secrets import UserSecretsClient\nuser_secrets = UserSecretsClient()\nuser_credential = user_secrets.get_gcloud_credential()\n\n# Set the credentials\nuser_secrets.set_tensorflow_credential(user_credential)\n\n# Use a familiar call to get the GCS path of the dataset\nfrom kaggle_datasets import KaggleDatasets\nGCS_DS_PATH = KaggleDatasets().get_gcs_path('hubmap-512x512-tfrecords-with-aug')\nGCS_DS_PATH","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_tf_gcs = GCS_DS_PATH+'/train/*.tfrecords'\nval_tf_gcs = GCS_DS_PATH+'/val/*.tfrecords'\ntrain_tf_files = tf.io.gfile.glob(train_tf_gcs)\nval_tf_files = tf.io.gfile.glob(val_tf_gcs)\nprint(val_tf_files[:3])\nprint(\"Train TFrecord Files:\", len(train_tf_files))\nprint(\"Val TFrecord Files:\", len(val_tf_files))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dataset = read_dataset(train_tf_files[15])\nvalidation_dataset = read_dataset(val_tf_files[15])\n\ntrain_image = []\ntrain_mask =[]\nfor image, mask in train_dataset.take(5):\n train_image, train_mask = image, mask\ntrain_mask = np.squeeze(train_mask)\n \ntest_image = []\ntest_mask =[]\nfor image, mask in validation_dataset.take(5):\n test_image, test_mask = image, mask\ntest_mask = np.squeeze(test_mask)\n \nfig, ax = plt.subplots(2,2,figsize=(20,10))\nax[0][0].imshow(train_image)\nax[0][1].imshow(train_mask)\nax[1][0].imshow(test_image)\nax[1][1].imshow(test_mask)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Model Building"},{"metadata":{"trusted":true},"cell_type":"code","source":"from tensorflow.keras.layers import *\nfrom tensorflow.keras.applications import *\nfrom tensorflow.keras.callbacks import *\nfrom tensorflow.keras.optimizers import Adam, Nadam\nfrom tensorflow.keras.metrics import *\nfrom tensorflow.keras.losses import binary_crossentropy\n\nnp.random.seed(13)\ntf.random.set_seed(13)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def squeeze_excite_block(inputs, ratio=8):\n init = inputs\n channel_axis = -1\n filters = init.shape[channel_axis]\n se_shape = (1, 1, filters)\n\n se = GlobalAveragePooling2D()(init)\n se = Reshape(se_shape)(se)\n se = Dense(filters // ratio, activation='relu', kernel_initializer='he_normal', use_bias=False)(se)\n se = Dense(filters, activation='sigmoid', kernel_initializer='he_normal', use_bias=False)(se)\n\n x = Multiply()([init, se])\n return x\n\ndef conv_block(inputs, filters):\n x = inputs\n\n x = Conv2D(filters, (3, 3), padding=\"same\")(x)\n x = BatchNormalization()(x)\n x = Activation('relu')(x)\n\n x = Conv2D(filters, (3, 3), padding=\"same\")(x)\n x = BatchNormalization()(x)\n x = Activation('relu')(x)\n\n x = squeeze_excite_block(x)\n\n return x\n\ndef encoder1(inputs):\n skip_connections = []\n\n model = VGG19(include_top=False, weights='imagenet', input_tensor=inputs)\n names = [\"block1_conv2\", \"block2_conv2\", \"block3_conv4\", \"block4_conv4\"]\n for name in names:\n skip_connections.append(model.get_layer(name).output)\n\n output = model.get_layer(\"block5_conv4\").output\n return output, skip_connections\n\ndef decoder1(inputs, skip_connections):\n num_filters = [256, 128, 64, 32]\n skip_connections.reverse()\n x = inputs\n shape = x.shape\n\n for i, f in enumerate(num_filters):\n x = Conv2DTranspose(shape[3], (2, 2), activation=\"relu\", strides=(2, 2))(x)\n x = Concatenate()([x, skip_connections[i]])\n x = conv_block(x, f)\n\n return x\n\ndef encoder2(inputs):\n num_filters = [32, 64, 128, 256]\n skip_connections = []\n x = inputs\n\n for i, f in enumerate(num_filters):\n x = conv_block(x, f)\n skip_connections.append(x)\n x = MaxPool2D((2, 2))(x)\n\n return x, skip_connections\n\ndef decoder2(inputs, skip_1, skip_2):\n num_filters = [256, 128, 64, 32]\n skip_2.reverse()\n x = inputs\n shape = x.shape\n\n for i, f in enumerate(num_filters):\n x = Conv2DTranspose(shape[3], (2, 2), activation=\"relu\", strides=(2, 2))(x)\n x = Concatenate()([x, skip_1[i], skip_2[i]])\n x = conv_block(x, f)\n\n return x\n\ndef output_block(inputs):\n x = Conv2D(1, (1, 1), padding=\"same\")(inputs)\n x = Activation('sigmoid')(x)\n return x\n\ndef ASPP(x, filter):\n shape = x.shape\n\n y1 = AveragePooling2D(pool_size=(shape[1], shape[2]))(x)\n y1 = Conv2D(filter, 1, padding=\"same\")(y1)\n y1 = BatchNormalization()(y1)\n y1 = Activation(\"relu\")(y1)\n shape2 = y1.shape\n \n y1 = Conv2DTranspose(shape2[3], (8,8), activation=\"relu\", strides=(shape[1], shape[2]))(y1)\n \n\n y2 = Conv2D(filter, 1, dilation_rate=1, padding=\"same\", use_bias=False)(x)\n y2 = BatchNormalization()(y2)\n y2 = Activation(\"relu\")(y2)\n\n y3 = Conv2D(filter, 3, dilation_rate=6, padding=\"same\", use_bias=False)(x)\n y3 = BatchNormalization()(y3)\n y3 = Activation(\"relu\")(y3)\n\n y4 = Conv2D(filter, 3, dilation_rate=12, padding=\"same\", use_bias=False)(x)\n y4 = BatchNormalization()(y4)\n y4 = Activation(\"relu\")(y4)\n\n y5 = Conv2D(filter, 3, dilation_rate=18, padding=\"same\", use_bias=False)(x)\n y5 = BatchNormalization()(y5)\n y5 = Activation(\"relu\")(y5)\n\n y = Concatenate()([y1, y2, y3, y4, y5])\n\n y = Conv2D(filter, 1, dilation_rate=1, padding=\"same\", use_bias=False)(y)\n y = BatchNormalization()(y)\n y = Activation(\"relu\")(y)\n\n return y\n\ndef build_model():\n inputs = Input((512, 512, 3))\n x, skip_1 = encoder1(inputs)\n x = ASPP(x, 64)\n x = decoder1(x, skip_1)\n outputs1 = output_block(x)\n\n x = inputs * outputs1\n\n x, skip_2 = encoder2(x)\n x = ASPP(x, 64)\n x = decoder2(x, skip_1, skip_2)\n outputs2 = output_block(x)\n outputs = Concatenate()([outputs1, outputs2])\n \n combine_output = Conv2D(1, (64, 64), activation=\"sigmoid\", padding=\"same\")(outputs)\n\n model = Model(inputs, combine_output)\n return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-output":true,"collapsed":true},"cell_type":"code","source":"model = build_model()\nmodel.summary(line_length=150)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Prepare Training Job Folder"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_job_path = '/kaggle/working/train_job'\nif not os.path.exists(train_job_path):\n os.makedirs(train_job_path)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from google.cloud import storage\nstorage_client = storage.Client(project='placesproject-284409')\n\nwith strategy.scope():\n\n def create_bucket(dataset_name):\n \"\"\"Creates a new bucket. https://cloud.google.com/storage/docs/ \"\"\"\n bucket = storage_client.create_bucket(dataset_name)\n print('Bucket {} created'.format(bucket.name))\n\n def upload_blob(bucket_name, source_file_name, destination_blob_name):\n \"\"\"Uploads a file to the bucket. https://cloud.google.com/storage/docs/ \"\"\"\n try:\n bucket = storage_client.get_bucket(bucket_name)\n blob = bucket.blob(destination_blob_name)\n blob.upload_from_filename(source_file_name)\n except Exception as E:\n print(\"Error uploading:\", E)\n # print('File {} uploaded to {}.'.format(\n # source_file_name,\n # destination_blob_name))\n\n def list_blobs(bucket_name):\n \"\"\"Lists all the blobs in the bucket. https://cloud.google.com/storage/docs/\"\"\"\n blob_list = []\n blobs = storage_client.list_blobs(bucket_name)\n for blob in blobs:\n blob_list.append(blob.name)\n #print(blob_list)\n return blob_list\n\n def download_to_kaggle(bucket_name,destination_directory,file_name):\n \"\"\"Takes the data from your GCS Bucket and puts it into the working directory of your Kaggle notebook\"\"\"\n os.makedirs(destination_directory, exist_ok = True)\n full_file_path = os.path.join(destination_directory, file_name)\n blobs = storage_client.list_blobs(bucket_name)\n for blob in blobs:\n blob.download_to_filename(full_file_path)","execution_count":null,"outputs":[]},{"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":"## Using GCS Bucket Persistant Job File Storage \n### Create a Bucket if it dosen't already exist"},{"metadata":{"trusted":true},"cell_type":"code","source":"bucket_name = 'hubmap_train_job_6' \ntry:\n create_bucket(bucket_name) \nexcept:\n pass","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Download the Model Weights and Callback Files\n\n### **If Model Weights and Callback Files exist from a previous session, download them to resume training from that point.**"},{"metadata":{"trusted":true},"cell_type":"code","source":"checkpoints_path = '/kaggle/working/train_job/checkpoints'\nif not os.path.exists(checkpoints_path):\n os.makedirs(checkpoints_path)\n \nfiles_list = list_blobs(bucket_name)\nckpt_list = [file for file in files_list if '.hdf5' in file]\n\nif ckpt_list:\n download_to_kaggle(bucket_name, train_job_path, 'lr_value.pickle')\n download_to_kaggle(bucket_name, checkpoints_path, ckpt_list[-1])\n checkpoint_path = os.path.join(checkpoints_path, ckpt_list[-1])\nelse:\n checkpoint_path = None","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Define Metrics"},{"metadata":{"trusted":true},"cell_type":"code","source":"with strategy.scope():\n def dice_coeff(y_true, y_pred):\n # add epsilon to avoid a divide by 0 error in case a slice has no pixels set\n # we only care about relative value, not absolute so this alteration doesn't matter\n _epsilon = 10 ** -7\n intersections = tf.reduce_sum(y_true * y_pred)\n unions = tf.reduce_sum(y_true + y_pred)\n dice_scores = (2.0 * intersections + _epsilon) / (unions + _epsilon)\n return dice_scores\n\n def dice_loss(y_true, y_pred):\n loss = 1 - dice_coeff(y_true, y_pred)\n return loss\n \n def bce_dice_loss(y_true, y_pred):\n return binary_crossentropy(y_true, y_pred) + dice_loss(y_true, y_pred)\n \n def tversky(y_true, y_pred, smooth=1, alpha=0.7):\n y_true_pos = tf.reshape(y_true,[-1])\n y_pred_pos = tf.reshape(y_pred,[-1])\n true_pos = tf.reduce_sum(y_true_pos * y_pred_pos)\n false_neg = tf.reduce_sum(y_true_pos * (1 - y_pred_pos))\n false_pos = tf.reduce_sum((1 - y_true_pos) * y_pred_pos)\n return (true_pos + smooth) / (true_pos + alpha * false_neg + (1 - alpha) * false_pos + smooth)\n\n def tversky_loss(y_true, y_pred):\n return 1 - tversky(y_true, y_pred)\n\n def focal_tversky_loss(y_true, y_pred, gamma=0.75):\n tv = tversky(y_true, y_pred)\n return K.pow((1 - tv), gamma)\n\n get_custom_objects().update({\"dice\": dice_loss})","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Functions to Load Records"},{"metadata":{"trusted":true},"cell_type":"code","source":"AUTO = tf.data.experimental.AUTOTUNE\nimage_feature_description = {\n 'height': tf.io.FixedLenFeature([], tf.int64),\n 'width': tf.io.FixedLenFeature([], tf.int64),\n 'num_channels': tf.io.FixedLenFeature([], tf.int64),\n 'img_bytes': tf.io.FixedLenFeature([], tf.string),\n 'mask': tf.io.FixedLenFeature([], tf.string),\n}\n\ndef _parse_image_function(example_proto):\n single_example = tf.io.parse_single_example(example_proto, image_feature_description)\n image = tf.reshape( tf.io.decode_raw(single_example['img_bytes'],out_type='uint8'), (512, 512, 3))\n mask = tf.reshape(tf.io.decode_raw(single_example['mask'],out_type='bool'),(512, 512, 1))\n ## normalize images array and cast image and mask to float32\n image = tf.cast(image, tf.float32) / 255.0\n mask = tf.cast(mask, tf.float32)\n return image, mask\n\ndef load_dataset(filenames, 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, compression_type=\"GZIP\")\n dataset = dataset.with_options(ignore_order)\n dataset = dataset.map(_parse_image_function, num_parallel_calls=AUTO)\n return dataset\n\ndef get_training_dataset():\n dataset = load_dataset(train_tf_files)\n #dataset = dataset.repeat()\n dataset = dataset.shuffle(47000)\n dataset = dataset.batch(16, drop_remainder=True)\n dataset = dataset.prefetch(AUTO)\n return dataset\ndef get_val_dataset():\n dataset = load_dataset(val_tf_files)\n #dataset = dataset.repeat()\n dataset = dataset.shuffle(5000)\n dataset = dataset.batch(16, drop_remainder=True)\n dataset = dataset.prefetch(AUTO)\n return dataset","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Custom ModelCheckpoint Funtion\n### Upload relevant weights and callback files to GCS for persistance\n\n\nWhen training resumes, the process may not start with the same conditions that took place when the checkpoint was saved. The learning rate would restart from its initial value as we may use learning rate decay or learning rate reduction on plateau. Thus, we ensure that here by pickling the learning rate at every epoch and persisting it in GCS."},{"metadata":{},"cell_type":"markdown","source":" Training Samples : 47703\n Validation Samples : 5829"},{"metadata":{"trusted":true},"cell_type":"code","source":"from keras.callbacks import ModelCheckpoint\nimport pickle\n\nwith strategy.scope():\n\n class ModelCheckpointEnhanced(ModelCheckpoint):\n def __init__(self, *args, **kwargs):\n # Added arguments\n self.lr_epoch_path = kwargs.pop('lr_epoch_path')\n \n super().__init__(*args, **kwargs)\n self.model_filepath = kwargs.pop('filepath')\n\n def on_epoch_end(self, epoch, logs=None):\n # Run normal flow:\n super().on_epoch_end(epoch,logs)\n\n model_filepath = self.model_filepath.format(epoch=epoch + 1, **logs)\n print(\"\\nModel_filepath\", model_filepath)\n lr_val = float(K.get_value(self.model.optimizer.lr))\n if self.epochs_since_last_save == 0 and epoch!=0:\n if self.save_best_only:\n current = logs.get(self.monitor)\n if current == self.best:\n # Note, there might be some cases where the last statement will save on unwanted epochs.\n # However, in the usual case where your monitoring value space is continuous this is not likely\n if os.path.exists(self.lr_epoch_path):\n os.remove(self.lr_epoch_path)\n \n with open(self.lr_epoch_path, \"wb\") as f:\n pickle.dump(lr_val, f)\n f.close()\n\n ## Uploading LR Pickle File to GCS\n file_name = os.path.basename(Path(self.lr_epoch_path))\n upload_blob(bucket_name, self.lr_epoch_path, file_name)\n \n ## Uploading Model Weight File to GCS\n model_file_name = os.path.basename(Path(model_filepath))\n upload_blob(bucket_name, model_filepath, model_file_name) \n else:\n if os.path.exists(self.lr_epoch_path):\n os.remove(self.lr_epoch_path)\n \n with open(self.lr_epoch_path, \"wb\") as f:\n pickle.dump(lr_val, f)\n f.close()\n \n ## Uploading LR Pickle File to GCS\n file_name = os.path.basename(Path(self.lr_epoch_path))\n upload_blob(bucket_name, self.lr_epoch_path, file_name)\n \n ## Uploading Model Weight File to GCS\n model_file_name = os.path.basename(Path(model_filepath))\n upload_blob(bucket_name, model_filepath, model_file_name)\n \n print('Model Files Uploaded to GCS')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"with strategy.scope(): \n def get_init_epoch(checkpoint_path):\n filename = os.path.basename(checkpoint_path)\n #filename = os.path.splitext(filename)[0]\n init_epoch = filename.split(\"_\")[1]\n print(\"init_epoch\", init_epoch)\n return int(init_epoch)\n \n metrics = [\n dice_coeff,\n bce_dice_loss,\n Recall(),\n Precision(),\n tversky_loss,\n focal_tversky_loss\n ]\n \n #Defining metrics as dict to pass to model load function\n metrics_dict = {\n 'dice_coeff': dice_coeff,\n 'bce_dice_loss': bce_dice_loss,\n 'tversky_loss': tversky_loss,\n 'focal_tversky_loss': focal_tversky_loss\n }\n \n \n ## Calling the Custom Checkpoint Callback.\n ## Passing the Weight path and the Learning Rate pickle file path which will be used resume training.\n ckpt_callback = ModelCheckpointEnhanced(filepath=train_job_path+'/checkpoints/weights_{epoch:02d}_{val_loss:.2f}.hdf5',\n monitor='val_loss', lr_epoch_path=train_job_path+'/lr_value.pickle')\n \n callbacks = [\n ckpt_callback,\n ReduceLROnPlateau(monitor='val_loss', factor=0.001, patience=3),\n EarlyStopping(monitor='val_loss', patience=3, restore_best_weights=False),\n CSVLogger(train_job_path+\"/data.csv\") \n ]\n \n train_dataset = get_training_dataset()\n validation_dataset = get_val_dataset()\n train_steps = round((47703//32)*0.3)\n validation_steps = round((5829//32)*0.3)\n# train_steps = 800\n# validation_steps = 100\n\n # Load checkpoint:\n if checkpoint_path is not None:\n # Load model:\n print(\"Resuming training from checkpoint:\", checkpoint_path)\n model = load_model(checkpoint_path, custom_objects=metrics_dict)\n \n # Finding the epoch index from which we are resuming\n initial_epoch = get_init_epoch(checkpoint_path)\n\n loaded_lr = pickle.load(open(train_job_path+'/lr_value.pickle', \"rb\" ))\n K.set_value(model.optimizer.lr, loaded_lr)\n \n else:\n print(\"Building model and starting training\")\n model = build_model()\n model.compile(optimizer = Adam(lr = 1e-2), loss = 'dice', metrics=metrics)\n initial_epoch = 0\n\n # Start/resume training\n model.fit(train_dataset, epochs=30, steps_per_epoch=train_steps,\n validation_data=validation_dataset, validation_steps=validation_steps,\n callbacks=callbacks,\n initial_epoch=initial_epoch)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"If the session stops, start a new session and just 'Run all' cells to resume training :)"}],"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}