{"cells":[{"metadata":{},"cell_type":"markdown","source":"**This notebook is Inference only using TPU. For training using custom loop with TPU please [see](http://www.kaggle.com/ashusma/training-lyft-tensorflow-tpu-multi-mode) this kernel.**\n\nModel has been trained on TPU with following parameters,\nimage size = [224, 224, 25]\n,epochs = 45\n,steps per epoch = 1000\n,batch_size = 192\naround 8.5 million samples"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import pandas as pd\nimport tensorflow as tf\nimport numpy as np\nfrom kaggle_datasets import KaggleDatasets\nfrom l5kit.evaluation import write_pred_csv\nimport time","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"try:\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":"sub = pd.read_csv('../input/lyft-motion-prediction-autonomous-vehicles/multi_mode_sample_submission.csv')\n\nTEST_DATA_DIR = 'lyft-test-tfrecords'\nTEST_GCS_PATH = KaggleDatasets().get_gcs_path(TEST_DATA_DIR)\ntest_files = tf.io.gfile.glob(TEST_GCS_PATH + '/test' +'/shard*.tfrecord')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"feature_descriptions = {\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 = {\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":"AUTO = tf.data.experimental.AUTOTUNE\nGLOBAL_BATCH_SIZE  = 64*strategy.num_replicas_in_sync\nIMG_SIZE = 224\nCHANNEL_DIM = 25\nTOTAL_STEPS = np.int(np.ceil(len(sub) / GLOBAL_BATCH_SIZE))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# we need to pad last partial batch with zeros to make each batch equal to global batch size\n# otherwise model will through an assertion error of shape mismatch on last batch because tf.distribute create dummy \n# variable of batch size 0 if it encounters any partial batch to make global batch divisible by no of replicas\n\ndef padded_batch(images, timestamp , track_id, world_from_agent, target, avail):\n    \n    if tf.shape(timestamp)[0] != GLOBAL_BATCH_SIZE:\n        images =    tf.concat([images , tf.zeros((GLOBAL_BATCH_SIZE - tf.shape(images)[0] , IMG_SIZE, IMG_SIZE, CHANNEL_DIM ))] , axis = 0)\n        timestamp = tf.concat([timestamp , tf.zeros((GLOBAL_BATCH_SIZE - tf.shape(timestamp)[0]) , tf.int64)] , axis = 0)\n        track_id =  tf.concat([track_id , tf.zeros((GLOBAL_BATCH_SIZE - tf.shape(track_id)[0]) , dtype = tf.int64)] , axis = 0)\n        world_from_agent = tf.concat([world_from_agent , tf.zeros((GLOBAL_BATCH_SIZE - tf.shape(world_from_agent)[0] , 3, 3) , tf.float64)] , axis = 0)\n    \n    return images , timestamp , track_id , world_from_agent , target , avail\n\ndef read_unlabeled_tfrecord(example):\n    \n    example = tf.io.parse_single_example(example, feature_descriptions)            # returns a dictionary\n    data = {k:tf.io.parse_tensor(example[k], feature_dtypes[k]) for k in example}\n    \n    image     =  tf.image.convert_image_dtype(data['image'], dtype = tf.float32)\n    image     =  tf.transpose(tf.squeeze(image), [1,2,0]) # converting images to format (batch_size, height, width, channel)\n    image     =  tf.reshape(image, (IMG_SIZE, IMG_SIZE, CHANNEL_DIM))\n    timestamp =  tf.squeeze(data['timestamp'])\n    track_id =   tf.squeeze(data['track_id'])\n    world_from_agent = tf.squeeze(data['world_from_agent'])\n    target =  tf.squeeze(data['target_positions'])\n    avail =  tf.squeeze(data['target_availabilities'])\n    return image , timestamp , track_id , world_from_agent , target, avail\n\ndef load_dataset(filenames, ordered = True, training = False):\n\n    # num_parallel reads is disabled to preserve the order of data\n    dataset = tf.data.TFRecordDataset(filenames,compression_type = 'GZIP' ) # tfrecords files are zipped\n    dataset = dataset.map(read_unlabeled_tfrecord, num_parallel_calls = AUTO).batch(GLOBAL_BATCH_SIZE,) \n    dataset = dataset.map(padded_batch , num_parallel_calls = AUTO) \n    return dataset","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_dataset(files ):\n    \n    dataset = load_dataset(files , training = False)\n    dataset = dataset.prefetch(AUTO)                   \n    \n     # each iteration of test_dataset contains data for all replicas which aggregates to global batch size\n    test_dataset = strategy.experimental_distribute_dataset(dataset) \n    return test_dataset","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\ndef transform_points(points, transf_matrix):\n    \n    # transform points 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)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def 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":{"trusted":true},"cell_type":"code","source":"class LyftModel(tf.keras.Model):\n    def __init__(self ,num_modes = 3,future_pred_frames = 50 ,):\n        super(LyftModel , self).__init__()\n        \n        self.conv1 = tf.keras.layers.Conv2D(3,kernel_size=1,use_bias=False,padding=\"same\" ,)\n        self.bn1   = tf.keras.layers.BatchNormalization()\n        self.relu  = tf.keras.layers.ReLU()\n        self.model = tf.keras.applications.ResNet50(include_top=False, weights='imagenet',)\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\n    def call(self, inputs):\n        x = self.conv1(inputs)\n        x = self.bn1(x)\n        x = self.relu(x)\n        x = self.model(x)\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\n        assert confidence.shape == (batch_size , self.num_modes) , f'confidence shape got {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        ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# keeping all model variables under the scope so that TPU can track them\n\ndef get_model():\n    with strategy.scope():\n#         tf.tpu.experimental.initialize_tpu_system()\n        model = LyftModel()\n        model.build((GLOBAL_BATCH_SIZE , None, None, CHANNEL_DIM))\n        model.summary()\n        model.load_weights('../input/lyft-resnet-model/epoch 45 and val_loss 23.7635 model.h5')\n        transf_points = lambda pred , world_from_agent : transform_points(pred, world_from_agent)\n        return model , transf_points","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# prediction function  \n# here tf.function compiles the function into tensorflow graph \n\n@tf.function\ndef test_step(image , world_from_agent):\n    pred , confidence = model(image , training = False)\n    pred = transf_points(pred , world_from_agent)\n    return  pred , confidence","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# prediction loop\n\ndef model_predict(test_dataset):\n    \n    start_time = epoch_start_time = time.time()\n\n    future_coord_offset = []\n    timestamps = []\n    track_ids = []\n    confs = []\n\n    print(f'predicting {TOTAL_STEPS} steps')\n    for step, (images, timestamp, track_id , world_from_agent , target , avail) ,in enumerate(test_dataset):\n         \n        # test dataset returns Per replica object dictionary whose keys contain index for each replicas and its values \n         #are dataset with batch_size GLOBAL_BATCH_SIZE / strategy.num_replicas_in_sync\n        \n        # in TPU : keys are  0, 1, 2, 3, 4, 5, 6, 7 for 8 replicas\n        \n        pred , confidence = strategy.run(test_step, args=(images, world_from_agent))\n        \n        global_batch_pred = np.concatenate([pred.values[i].numpy() for i in range(strategy.num_replicas_in_sync)], axis = 0)\n        global_batch_confs = np.concatenate([confidence.values[i].numpy() for i in range(strategy.num_replicas_in_sync)], axis = 0)\n        global_batch_timestamps = np.concatenate([timestamp.values[i].numpy() for i in range(strategy.num_replicas_in_sync)], axis = 0)\n        global_batch_track_ids = np.concatenate([track_id.values[i].numpy() for i in range(strategy.num_replicas_in_sync)], axis = 0 )\n        \n        future_coord_offset.append(global_batch_pred)\n        confs.append(global_batch_confs)\n        timestamps.append(global_batch_timestamps)\n        track_ids.append(global_batch_track_ids)\n         \n        print('=' , end = ' ' , flush = True)\n    epoch_time =   time.time() - epoch_start_time\n    print('time: {:0.1f}s'.format(epoch_time))\n    \n    return future_coord_offset, timestamps , track_ids, confs","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# now getting the dataset and model\n\ntest_dataset = get_dataset(test_files)\nmodel , transf_points = get_model()\nfuture_coord_offset, timestamps , track_ids , confs = model_predict(test_dataset )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# making csv file using l5kit built in function\n# since last batch contains zeros that we padded so we will discard them by taking only length of test set\n\npred_path = 'submission.csv'\ntest_length = len(sub)\nwrite_pred_csv(pred_path,\n              timestamps =  np.concatenate(timestamps, axis = 0)[0:test_length], \n              track_ids  =  np.concatenate(track_ids , axis = 0)[0:test_length],   \n              coords     =  np.concatenate(future_coord_offset , axis = 0)[0:test_length],   \n              confs      =  np.concatenate(confs , axis = 0)[0:test_length],   \n              )","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}