{"cells":[{"metadata":{},"cell_type":"markdown","source":"# **This notebook demonstrates the training of LYFT data on Tensorflow TPU using custom training loop .** \n\nFew keypoints of the notebook are :\n\n\n* Since LYFT data is very large so i have zipped the tfrecord files to store more data into one file, tfrecord files contains images of size 224 and channel dimension 25 , num_history_channel 10 . I have added the target positions and availabilities to the validation tfrecord just like training tfrecord that will allow evaluation during training and also faster experiments.\n\n* To increase the data input to TPU vectorization(operate on batches) on user defined function has been done before map function that increases the throughput upto 4 times and significantly reduces the training time .To make this happen tf.io.parse_example is used instead of tf.io.parse_single_example that parses single example at a time\n\n* dataset.batch(drop remainder = True) is used for both training and validation that is slighlty faster as discussed [here](http://https://www.kaggle.com/mgornergoogle/custom-training-loop-with-100-flowers-on-tpu), Although it will drop some validation example but its ok as we have large amount of validation samples.  \n\n* For validation dataset.cache() has not been used due to the large size of data that results in memory overflow and TPU throws socket closed error after few epochs.\n\n* Custom loss ,Transform points function has been modified to calculate on batches. \n\nall the other necessary information have been given under comments in the code"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"\nimport tensorflow as tf\nimport pandas as pd\nimport numpy as np\nfrom pathlib import Path\nfrom kaggle_datasets import KaggleDatasets\nfrom collections import namedtuple\nimport matplotlib.pyplot as plt\nimport time\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# TPU Detection And Initialization"},{"metadata":{"trusted":true},"cell_type":"code","source":"\ntry:\n    # TPU detection. No parameters necessary if TPU_NAME environment variable is\n    # set: this is always the case on Kaggle.\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    # Default distribution strategy in Tensorflow. Works on CPU and single GPU.\n    strategy = tf.distribute.get_strategy()\n\nprint(\"REPLICAS: \", strategy.num_replicas_in_sync)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# combining all the tfrecord files\n\ntrain_files = []\nfor i in range(3):\n    TRAIN_GCS_PATH = KaggleDatasets().get_gcs_path(f'lyft-training-tfrecord-224-{i}')\n    train_files += tf.io.gfile.glob(TRAIN_GCS_PATH + '/training' + '/shard*.tfrecord' )\n \n\n\nVALID_GCS_PATH = KaggleDatasets().get_gcs_path('lyft-validation-tfrecord-224')\nvalid_files = tf.io.gfile.glob(VALID_GCS_PATH  +'/validation' + '/shard*.tfrecord' )\n\nnp.random.shuffle(train_files)       ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"AUTO = tf.data.experimental.AUTOTUNE\nIMG_DIM = 224\nCHANNEL_DIM = 25\nEPOCHS = 10\nSTEPS_PER_EPOCH = 500\nVALIDATION_STEP = 200\nGLOBAL_BATCH_SIZE = 16*strategy.num_replicas_in_sync ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataset PipeLine"},{"metadata":{"trusted":true},"cell_type":"code","source":"\nfeature_descriptions_train = {\n    'image': tf.io.FixedLenFeature([], tf.string),\n    'target_positions': tf.io.FixedLenFeature([], tf.string),\n    'target_availabilities': tf.io.FixedLenFeature([], tf.string),\n}\n\nfeature_dtypes_train = {\n    'image': tf.uint8,\n    'target_positions': tf.float32,\n    'target_availabilities': tf.uint8,\n}\n\nfeature_descriptions_valid = {\n    'image': tf.io.FixedLenFeature([], tf.string),\n    'target_positions': tf.io.FixedLenFeature([], tf.string),\n    'target_availabilities': tf.io.FixedLenFeature([], tf.string),\n    'target_yaws': tf.io.FixedLenFeature([], tf.string),\n    'world_from_agent' : tf.io.FixedLenFeature([], tf.string),\n    'history_positions': tf.io.FixedLenFeature([], tf.string),\n    'history_yaws': tf.io.FixedLenFeature([], tf.string),\n    'history_availabilities': tf.io.FixedLenFeature([], tf.string),\n    'world_to_image': tf.io.FixedLenFeature([], tf.string),\n    'track_id': tf.io.FixedLenFeature([], tf.string),\n    'timestamp': tf.io.FixedLenFeature([], tf.string),\n    'centroid': tf.io.FixedLenFeature([], tf.string),\n    'yaw': tf.io.FixedLenFeature([], tf.string),\n    'extent': tf.io.FixedLenFeature([], tf.string),\n}\n\nfeature_dtypes_valid = {\n    'image': tf.uint8,\n    'target_positions': tf.float32,\n    'target_availabilities': tf.float32,\n    'target_yaws': tf.float32,\n    'world_from_agent' : tf.float64,\n    'history_positions': tf.float32,\n    'history_yaws':tf.float32,\n    'history_availabilities': tf.float32,\n    'world_to_image': tf.float64,\n    'track_id': tf.int64,\n    'timestamp': tf.int64,\n    'centroid': tf.float64,\n    'yaw': tf.float64,\n    'extent': tf.float32,\n}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def read_train_tfrecord(example):\n    \n    # tf.io.parse_example parse examples in batches\n    example = tf.io.parse_example(example, feature_descriptions_train)\n    data = {}\n    data['image']  =   [tf.squeeze(tf.io.parse_tensor(example['image'][i], tf.uint8)) for i in range(GLOBAL_BATCH_SIZE)]\n    data['target_availabilities'] = [tf.squeeze(tf.io.parse_tensor(example['target_availabilities'][i]  , tf.uint8)) for i in range(GLOBAL_BATCH_SIZE)]\n    data['target_positions']  = [tf.squeeze(tf.io.parse_tensor(example['target_positions'][i], tf.float32)) for i in range(GLOBAL_BATCH_SIZE)]\n\n    image = tf.image.convert_image_dtype(data['image'], dtype = tf.float32)\n    target_avail  = tf.image.convert_image_dtype(data['target_availabilities'], dtype = tf.float32)\n    \n    image = tf.transpose(image , [0, 2, 3, 1])               # converting images to format (batch_size, height, width, channel)\n    image = tf.reshape(image, shape = (GLOBAL_BATCH_SIZE,IMG_DIM, IMG_DIM, CHANNEL_DIM))\n    target_pos = tf.reshape(data['target_positions'] , shape = (GLOBAL_BATCH_SIZE, 50, 2))\n    target_avail =  tf.reshape(target_avail , shape = (GLOBAL_BATCH_SIZE, 50))\n    return image , target_pos , target_avail\n\n\n\ndef read_validation_tfrecord(example):\n    \n    example = tf.io.parse_example(example, feature_descriptions_valid)\n    data = {}\n    data['image'] = [tf.io.parse_tensor(example['image'][i], tf.uint8) for i in range(GLOBAL_BATCH_SIZE)],\n    data['target_availabilities'] = [tf.io.parse_tensor(example['target_availabilities'][i], tf.float32) for i in range(GLOBAL_BATCH_SIZE)]\n    data['target_positions']  = [tf.io.parse_tensor(example['target_positions'][i], tf.float32) for i in range(GLOBAL_BATCH_SIZE)],\n    data['world_from_agent'] = [tf.io.parse_tensor(example['world_from_agent'][i], tf.float64) for i in range(GLOBAL_BATCH_SIZE)],\n    \n    image = tf.image.convert_image_dtype(data['image'], dtype = tf.float32)\n    \n    image = tf.transpose(tf.squeeze(image), [0,2,3,1]) \n    image = tf.reshape(image, shape = (GLOBAL_BATCH_SIZE,IMG_DIM, IMG_DIM, CHANNEL_DIM)) \n    target_pos = tf.reshape(tf.squeeze(data['target_positions']) , shape = (GLOBAL_BATCH_SIZE, 50, 2))\n    target_avail = tf.reshape(tf.squeeze(data['target_availabilities']) , shape = (GLOBAL_BATCH_SIZE, 50))\n    world_from_agent = tf.reshape(tf.squeeze(data['world_from_agent']) , shape = (GLOBAL_BATCH_SIZE ,3 , 3))\n    return image , target_pos , target_avail , world_from_agent","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def load_dataset(filenames, labeled = True, ordered = False , training = True):\n    # Read from TFRecords. For optimal performance, reading from multiple files at once and\n    # Diregarding data order. Order does not matter since we will be shuffling the data anyway\n    \n    ignore_order = tf.data.Options()\n    if not ordered:\n        # disable order, increase speed\n        ignore_order.experimental_deterministic = False \n        \n    # automatically interleaves reads from multiple files\n    # tfrecords file are zipped\n    dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads = AUTO , compression_type = 'GZIP')\n\n    # use data as soon as it streams in, rather than in its original order\n    dataset = dataset.with_options(ignore_order)\n    \n    # here we are vectorizing(operate over batches) the function by doing batch transformation before map function\n    dataset = dataset.batch(GLOBAL_BATCH_SIZE , drop_remainder = True)\n    dataset = dataset.map(read_train_tfrecord if training else read_validation_tfrecord, num_parallel_calls = AUTO)\n    return dataset\n\n\ndef get_dataset(files , training = True):\n   \n    if training:\n        dataset = load_dataset(files , training = True)\n        dataset = dataset.shuffle(128).prefetch(AUTO)\n\n    else:\n        # we are not caching the validation data due to larger size which will result in memory overflow\n        dataset = load_dataset(files , training = False)\n        dataset = dataset.prefetch(AUTO)\n        \n    return dataset","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Visualization\nNow we will visualize the model input\n"},{"metadata":{"trusted":true},"cell_type":"code","source":"\n# loading the dataset\ntraining_dataset = get_dataset(train_files, training = True)\nvalidation_dataset = get_dataset(valid_files , training = False)\nimage , target_pos , target_avail = next(iter(training_dataset))\nprint('shapes of images, target_positions, target_avail are ', (image.shape , target_pos.shape , target_avail.shape))\n\n#plotting input image\nplt.figure(figsize=(15, 15))\nfor i in range(25):\n    plt.subplot(5, 5, i+1)\n    plt.imshow(image[45][:, :, i])\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Custom Metric"},{"metadata":{"trusted":true},"cell_type":"code","source":"def neg_multi_log_likelihood(gt, pred, confidences, avails,):\n    \n    # function calculates loss in Batch\n    \n    assert len(pred.shape) == 4, f\"expected 4D (B,M,T,C) array for pred, got {pred.shape}\"\n    batch_size , num_modes, future_len, num_coords = pred.shape\n\n    assert gt.shape == (batch_size , future_len, num_coords), f\"expected 3D (Batch, Time , Coords) array for gt, got {gt.shape}\"\n    assert confidences.shape == (batch_size , num_modes,), f\"expected 2D (Batch, Modes) array for gt, got {confidences.shape}\"\n    assert avails.shape == (batch_size , future_len,), f\"expected 2D (Batch,Time) array for gt, got {avails.shape}\"\n    \n    gt = tf.expand_dims(gt, axis = 1)  # add modes\n    avails = avails[: , None, :, None]  # add modes and cords\n    error = tf.math.reduce_sum(((gt - pred) * avails) ** 2, axis=-1)    # reduce coords and use availability\n    confidences = tf.clip_by_value(confidences , clip_value_min = tf.pow(0.1, 12) , clip_value_max = 1.0)   # to avoid exploding gradient\n    error = tf.math.log(confidences ) - 0.5 * tf.math.reduce_sum(error, axis=-1)  # reduce time\n\n    # use max aggregator on modes for numerical stability\n    max_value = tf.math.reduce_max(error, axis=1, keepdims=True)  # error are negative at this point, so max() gives the minimum one\n    error = -tf.math.log(tf.math.reduce_sum(tf.exp(error - max_value), axis=-1, keepdims=True)) - max_value  # reduce modes\n    return error\n\n\n   \ndef transform_points(points, transf_matrix):\n    \n    # transform prediction to world coordinates in batches\n    \n    transf_matrix = tf.expand_dims(transf_matrix , axis = -1)\n    assert len(points.shape) == len(transf_matrix.shape) == 4, (\n    f\"dimensions mismatch, both points ({points.shape}) and \"\n    f\"transf_matrix ({transf_matrix.shape}) needs to be tensors of rank 4.\"\n    )\n\n    if points.shape[3] not in [2, 3]:\n        raise AssertionError(f\"Points input should be (N, 2) or (N, 3) shape, received {points.shape}\")\n\n    assert points.shape[3] == transf_matrix.shape[2] - 1, \"points dim should be one less than matrix dim\"\n\n    points = tf.cast(points , tf.float64)\n    points = tf.matmul(points , tf.transpose(transf_matrix[:, :-1, :-1, :] , perm = [0,3,2,1])) \n    return tf.cast(points , tf.float32)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# changing input channel to match rasterizer output will change the weights dimension of the layer that will \n# take input , so accordingly we need to change the weights dimension of that layer\n\n# resnet50 first convolution layer has weights shape (7, 7, 3, 64) for 3 input channel\n# and required is (7, 7, 25, 64) for 25 input channel\n\ndef modified_resnet50():\n    \n     # model with 3 input channel dim with pretrained weights\n    pretrained_model = tf.keras.applications.ResNet50(include_top=False, weights='imagenet', input_shape = (None, None, 3)) \n    \n    # model with 25 input channel dim without pretrained weight\n    modified_model = tf.keras.applications.ResNet50(include_top=False, weights= None, input_shape = (None, None, 25))                \n\n    for pretrained_model_layer , modified_model_layer in zip(pretrained_model.layers , modified_model.layers):\n        layer_to_modify = ['conv1_conv']                    # conv1_conv is name of layer that takes the input and will be modified\n        if pretrained_model_layer.name in layer_to_modify :          \n            kernel = pretrained_model_layer.get_weights()[0]  # kernel weight shape is (7, 7 ,3, 64)\n            bias = pretrained_model_layer.get_weights()[1]\n            \n            # concatenating along channel axis to make channel dimension 25\n            weights = np.concatenate((kernel[:, :, -1: ,:] , np.tile( kernel , [1, 1, 8, 1]) ) , axis=  -2)  \n            modified_model_layer.set_weights((weights , bias))\n        else:\n            modified_model_layer.set_weights(pretrained_model_layer.get_weights())\n\n    return modified_model","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Define Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"\nclass LyftModel(tf.keras.Model):\n    def __init__(self ,num_modes = 3,future_pred_frames = 50 ,):\n        super(LyftModel , self).__init__()\n        \n        self.model = modified_resnet50()\n        self.gap = tf.keras.layers.GlobalAveragePooling2D()\n        self.dropout = tf.keras.layers.Dropout(0.2)\n        \n        self.future_len = num_modes * future_pred_frames * 2      \n        self.future_pred_frames = future_pred_frames\n        self.num_modes = num_modes\n        \n        self.dense1 = tf.keras.layers.Dense(self.future_len + self.num_modes ,)\n  \n\n        \n    def call(self,inputs):\n        x = self.model(inputs)\n        x = self.gap(x)\n        x = self.dropout(x)\n        x  = self.dense1(x)\n        \n        batch_size, _  = x.shape\n        pred , confidence = tf.split(x , num_or_size_splits = [self.future_len, self.num_modes], axis = 1)\n        assert confidence.shape == (batch_size , self.num_modes) , f'confidence got shape {confidence.shape}'\n        pred = tf.reshape(pred , shape = (batch_size ,self.num_modes, self.future_pred_frames , 2))\n        confidence = tf.nn.softmax(confidence , axis = 1)\n        return pred , confidence \n        \n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# defining the model and variables\n\ndef get_model_and_variables():\n    model = LyftModel()\n    model.build((GLOBAL_BATCH_SIZE , IMG_DIM , IMG_DIM , CHANNEL_DIM))\n    model.summary()\n    optimizer = tf.keras.optimizers.Adam(learning_rate= 0.001)\n\n    # tf.nn.compute_average_loss will aggregate the per example loss across all replicas and returns the average scalar loss \n    loss_func = lambda a, b, c, d : tf.nn.compute_average_loss(neg_multi_log_likelihood(a, b, c, d) , \n                                global_batch_size = GLOBAL_BATCH_SIZE )\n    \n    transf_points = lambda pred , world_from_agent : transform_points(pred, world_from_agent)\n\n    # metrics\n    training_loss = tf.keras.metrics.Sum()       \n    validation_loss = tf.keras.metrics.Sum()\n    return model , optimizer , loss_func , transf_points ,training_loss , validation_loss\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train and Validation step function"},{"metadata":{"trusted":true},"cell_type":"code","source":"# @tf.function compiles the function to tensorflow graph to run on tpu\n\n@tf.function\ndef train_step(image, target_pos , target_avail): \n    # this function will run on each replica and constitutes one training step\n    \n    with tf.GradientTape() as tape :\n        pred , confidence = model(image , training = True)\n        loss_value = loss_func(target_pos , pred, confidence , target_avail)\n    grads = tape.gradient(loss_value , model.trainable_variables)\n    optimizer.apply_gradients(list(zip(grads , model.trainable_variables)))\n    training_loss.update_state(loss_value)\n\n        \n@tf.function\ndef valid_step(image , target_pos, target_avail, world_from_agent):\n    pred , confidence = model(image , training = False)\n    pred = transf_points(pred , world_from_agent)          # transforming the points to world space coordinates\n    val_loss = loss_func(target_pos , pred, confidence , target_avail)  # calculating validation loss\n    validation_loss.update_state(val_loss)\n\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Custom Training Loop"},{"metadata":{"trusted":true},"cell_type":"code","source":"# training loop\n\ndef train_model(train_dataset , valid_dataset ,EPOCHS =  5, STEPS_PER_EPOCH = 200 ,VALIDATION_STEP = 100):\n    # now we will distribute the dataset according to the strategy here it is TPUStrategy\n    \n    train_dist_dataset = strategy.experimental_distribute_dataset(train_dataset)  \n    valid_dist_dataset = strategy.experimental_distribute_dataset(valid_dataset)\n    \n    start_time = epoch_start_time = time.time()\n\n    print(\"Steps per epoch:\", STEPS_PER_EPOCH)\n    History = namedtuple('History', 'history')\n    history = History(history={'train_loss': [], 'val_loss': [],})\n\n    epoch = 0\n    for step , (images, target_pos, target_avails) in enumerate(train_dist_dataset):\n\n        # each iteration on train dist dataset returns per replica object dictionary containing data for each worker or replica\n        # batch size for each replica is GLOBAL_BATCH_SIZE /strategy.num_replicas_in_sync\n        \n        #strategy.run will distribute train_step and execute operation as specified by function on each replica \n\n        strategy.run(train_step, args=(images, target_pos, target_avails))\n        #print('=' , end = ' ' , flush = True)\n\n        # validation run at the end of each epoch\n        if ((step+1) // STEPS_PER_EPOCH) > epoch:\n\n            # validation run\n            for val_step ,(images , target_pos , target_avails , world_from_agent) in enumerate(valid_dist_dataset):\n                strategy.run(valid_step, args=(images, target_pos, target_avails ,world_from_agent))\n                if (val_step + 1) % VALIDATION_STEP == 0: \n                    break\n\n            # storing result \n            history.history['train_loss'].append(training_loss.result().numpy() / STEPS_PER_EPOCH )\n            history.history['val_loss'].append((validation_loss.result().numpy() / VALIDATION_STEP))\n\n            # show metrics\n            epoch_time = time.time() - epoch_start_time\n            print('\\nEPOCH {:d}/{:d}'.format(epoch+1, EPOCHS))\n            print('time: {:0.1f}s'.format(epoch_time),\n                  'loss: {:0.4f}'.format(history.history['train_loss'][-1]),\n                  'val_loss: {:0.4f}'.format(history.history['val_loss'][-1]),)\n            \n            # saving the model\n            model.save_weights('./epoch {:d}, train_loss {:0.4f}, val_loss {:0.4f} model.h5'.format(epoch+1 , \n                                                                               history.history['train_loss'][-1],\n                                                                               history.history['val_loss'][-1]  ))                   \n                                                                              \n            \n            # set up next epoch\n            \n            epoch = (step+1) // STEPS_PER_EPOCH\n            epoch_start_time = time.time()\n            validation_loss.reset_states()\n            training_loss.reset_states()\n        \n        if epoch >= EPOCHS:\n            break\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# calling the model and variables under strategy.scope to allow tpu to track and compute the variables\n\nwith strategy.scope():\n    model , optimizer , loss_func , transf_points ,training_loss , validation_loss = get_model_and_variables()\n    \n\ntrain_model(training_dataset ,validation_dataset ,EPOCHS = EPOCHS ,STEPS_PER_EPOCH = STEPS_PER_EPOCH ,VALIDATION_STEP = VALIDATION_STEP)","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}