{"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":"## Introduction\n\nThis notebook is based on the notebook originally published with LSTM training and inference: [3 LSTMs; with Data Picking and Shifting](https://www.kaggle.com/code/seungmoklee/3-lstms-with-data-picking-and-shifting). So if you like my notebooks don't forget the work that it is based on!\n\nThis notebook shows how the LSTM models can be trained on TPU. It should also run on GPU...just take into account the smaller batch size and learning rate in that case...\n\nThe datapreprocessing is done in the same way with the following exceptions:\n1. The pulse_count is 96 instead of 128. During various experiments I often noticed that the performance was better with a maximum of 96 pulses compared to 128 pulses.\n2. The features 'r_err' and 'z_err' where not added. In various experiments with the model training they seemed to hurt the performance.\n\nWith these modifications the files are a lot smaller and more files will fit into memory. In the attached Dataset 90 training files are provided. See the following notebook for how to preprocess the training data and generate the required files: [Tensorflow LSTM Model Data PreProcessor](https://www.kaggle.com/rsmits/tensorflow-lstm-model-data-preprocessor)\n\nI made the following modifications that (all combined..) increased the performance of the LSTM model drastically. Below the most important ones:\n* Use a lot more of the training data. The largest set I used sofar is 70 training files.\n* Increase the bin_num. Further increasing might be possible... I did notice that this only works when enough training files are used. With only a small set of training files performance will actually drop.\n* Use a large batch size\n* Higher learning rate\n* Use GRU instead of LSTM\n* Use a larger number of LSTM cells.\n* Add an additional Bidirectional/GRU layer.\n* Add an additional Dense layer with enough units.\n* Add a Masking layer. This takes into account any pulses with all 0. values.\n* Don't use One Hot Encoding, just an integer. One Hot Encoding is killing the amount of RAM available.\n* Lower the maximum pulse count from 128 to 96.\n* Use only 6 features when training: time, charge, aux, x, y, z \n\nThe model files used in the [Tensorflow LSTM Model Inference](https://www.kaggle.com/code/rsmits/tensorflow-lstm-model-inference) notebook where trained on my local laptop loading in 70 training files into 32 GB of RAM. All code in this TPU training notebook is the same with the difference that every (data_new_load_interval) number of epochs a new set of training files (amount set with train_files_delta) is loaded. This takes into account the limited RAM available in the Kaggle environment while still being able to use all data...just in a different way.\n\nTo get the same performance as the model files in the inference notebook you should train locally with all data loaded in RAM. The batch loading of training files as used for TPU does not seem to achieve the exact high performance. With only 20 hours of TPU time available each week it is a limitation to be able to tweak and tune the TPU notebook to do exactly the same. To mimic the local training from the 2 inference models change the hyperparameters as they are marked with 'Local Training'.\n\nSome ideas for increasing the performance:\n* Change the current .npz files to TFRecords. That will allow all data to be used easily by the TPU's.\n* Experiment with different bin_nums.\n* Experiment with Dropout.\n* Experiment with Learning Rate scheduling.\n* Etc.\n\nLet me know if you have any feedback, comments or questions about this notebook. ","metadata":{}},{"cell_type":"markdown","source":"## Update Latest Version\n\nIn this latest update I changed a few hyperparameters to increase the performance when training on TPU. The delta set of training files used is now reloaded more often and the selection of training files is completely random.\n\nAfter some verification it turns out that using Generators with TPU is not supported. As already mentioned in the earlier versions of this training notebook you can further experiment with the notebook as it currently is ... or switch to using TFRecords. Tate Larkin made a very nice [notebook](https://www.kaggle.com/code/tatelarkin/saving-and-loading-icecube-data-as-tfrecord) showing the concepts of how to convert the data into TFRecords.","metadata":{}},{"cell_type":"code","source":"# Import\nimport numpy as np\nimport os\nimport gc\nimport tensorflow as tf\nimport random\nfrom tqdm.notebook import tqdm","metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","papermill":{"duration":0.025652,"end_time":"2023-02-27T22:23:47.490482","exception":false,"start_time":"2023-02-27T22:23:47.46483","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Configure Strategy. Assume TPU...if not set default for GPU\ntpu = None\ntry:\n    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n    tf.config.experimental_connect_to_cluster(tpu)\n    tf.tpu.experimental.initialize_tpu_system(tpu)\n    strategy = tf.distribute.experimental.TPUStrategy(tpu)\nexcept:\n    strategy = tf.distribute.get_strategy()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Constants","metadata":{"papermill":{"duration":0.00737,"end_time":"2023-02-27T22:23:47.505631","exception":false,"start_time":"2023-02-27T22:23:47.498261","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# Training\nvalidation_files_amount = 1\ndata_new_load_interval = 6      # Local Training: None\ntrain_files_delta = 15          # Local Training: None\nepochs = 75                     # Local Training: 30\nbatch_size = 8192               # Local Training: 2048\nlearning_rate = 0.0022          # Local Training: 0.0005\nverbose = 0\n\n# Training Batches\ntrain_batch_id_min = 100\ntrain_batch_id_max = 190\ntrain_batch_ids = [*range(train_batch_id_min, train_batch_id_max+1)]\nnp.random.shuffle(train_batch_ids)\nprint(train_batch_ids)\n\n# Model Parameters\npulse_count = 96\nfeature_count = 6\nlstm_units = 192\nbin_num = 24\n\n# Data\nbase_dir = \"/kaggle/input/lstmicecubesdata/\"\nfile_format = base_dir + 'pp_mpc96_n7_batch_{batch_id:d}.npz'","metadata":{"papermill":{"duration":0.016643,"end_time":"2023-02-27T22:23:47.529892","exception":false,"start_time":"2023-02-27T22:23:47.513249","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Set Seed\nseed = 4242\ntf.random.set_seed(seed)\nrandom.seed(seed)\nnp.random.seed(seed)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Prepare Metric","metadata":{"papermill":{"duration":0.007186,"end_time":"2023-02-27T22:23:47.59189","exception":false,"start_time":"2023-02-27T22:23:47.584704","status":"completed"},"tags":[]}},{"cell_type":"code","source":"def angular_dist_score(az_true, zen_true, az_pred, zen_pred):\n    '''\n    calculate the MAE of the angular distance between two directions.\n    The two vectors are first converted to cartesian unit vectors,\n    and then their scalar product is computed, which is equal to\n    the cosine of the angle between the two vectors. The inverse \n    cosine (arccos) thereof is then the angle between the two input vectors\n    \n    Parameters:\n    -----------\n    \n    az_true : float (or array thereof)\n        true azimuth value(s) in radian\n    zen_true : float (or array thereof)\n        true zenith value(s) in radian\n    az_pred : float (or array thereof)\n        predicted azimuth value(s) in radian\n    zen_pred : float (or array thereof)\n        predicted zenith value(s) in radian\n    \n    Returns:\n    --------\n    \n    dist : float\n        mean over the angular distance(s) in radian\n    '''\n    \n    if not (np.all(np.isfinite(az_true)) and\n            np.all(np.isfinite(zen_true)) and\n            np.all(np.isfinite(az_pred)) and\n            np.all(np.isfinite(zen_pred))):\n        raise ValueError(\"All arguments must be finite\")\n    \n    # pre-compute all sine and cosine values\n    sa1 = np.sin(az_true)\n    ca1 = np.cos(az_true)\n    sz1 = np.sin(zen_true)\n    cz1 = np.cos(zen_true)\n    \n    sa2 = np.sin(az_pred)\n    ca2 = np.cos(az_pred)\n    sz2 = np.sin(zen_pred)\n    cz2 = np.cos(zen_pred)\n    \n    # scalar product of the two cartesian vectors (x = sz*ca, y = sz*sa, z = cz)\n    scalar_prod = sz1*sz2*(ca1*ca2 + sa1*sa2) + (cz1*cz2)\n    \n    # scalar product of two unit vectors is always between -1 and 1, this is against nummerical instability\n    # that might otherwise occure from the finite precision of the sine and cosine functions\n    scalar_prod =  np.clip(scalar_prod, -1, 1)\n    \n    # convert back to an angle (in radian)\n    return np.average(np.abs(np.arccos(scalar_prod)))","metadata":{"papermill":{"duration":0.021907,"end_time":"2023-02-27T22:23:47.637698","exception":false,"start_time":"2023-02-27T22:23:47.615791","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Define Azimuth and Zenith Bins","metadata":{}},{"cell_type":"code","source":"# Create Azimuth Edges\nazimuth_edges = np.linspace(0, 2 * np.pi, bin_num + 1)\nprint(azimuth_edges)\n\n# Create Zenith Edges\nzenith_edges = []\nzenith_edges.append(0)\nfor bin_idx in range(1, bin_num):\n    zenith_edges.append(np.arccos(np.cos(zenith_edges[-1]) - 2 / (bin_num)))\nzenith_edges.append(np.pi)\nzenith_edges = np.array(zenith_edges)\nprint(zenith_edges)","metadata":{"papermill":{"duration":0.019686,"end_time":"2023-02-27T22:23:47.680761","exception":false,"start_time":"2023-02-27T22:23:47.661075","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Supporting Functions","metadata":{}},{"cell_type":"code","source":"angle_bin_zenith0 = np.tile(zenith_edges[:-1], bin_num)\nangle_bin_zenith1 = np.tile(zenith_edges[1:], bin_num)\nangle_bin_azimuth0 = np.repeat(azimuth_edges[:-1], bin_num)\nangle_bin_azimuth1 = np.repeat(azimuth_edges[1:], bin_num)\n\nangle_bin_area = (angle_bin_azimuth1 - angle_bin_azimuth0) * (np.cos(angle_bin_zenith0) - np.cos(angle_bin_zenith1))\nangle_bin_vector_sum_x = (np.sin(angle_bin_azimuth1) - np.sin(angle_bin_azimuth0)) * ((angle_bin_zenith1 - angle_bin_zenith0) / 2 - (np.sin(2 * angle_bin_zenith1) - np.sin(2 * angle_bin_zenith0)) / 4)\nangle_bin_vector_sum_y = (np.cos(angle_bin_azimuth0) - np.cos(angle_bin_azimuth1)) * ((angle_bin_zenith1 - angle_bin_zenith0) / 2 - (np.sin(2 * angle_bin_zenith1) - np.sin(2 * angle_bin_zenith0)) / 4)\nangle_bin_vector_sum_z = (angle_bin_azimuth1 - angle_bin_azimuth0) * ((np.cos(2 * angle_bin_zenith0) - np.cos(2 * angle_bin_zenith1)) / 4)\n\nangle_bin_vector_mean_x = angle_bin_vector_sum_x / angle_bin_area\nangle_bin_vector_mean_y = angle_bin_vector_sum_y / angle_bin_area\nangle_bin_vector_mean_z = angle_bin_vector_sum_z / angle_bin_area\n\nangle_bin_vector = np.zeros((1, bin_num * bin_num, 3))\nangle_bin_vector[:, :, 0] = angle_bin_vector_mean_x\nangle_bin_vector[:, :, 1] = angle_bin_vector_mean_y\nangle_bin_vector[:, :, 2] = angle_bin_vector_mean_z\n\ndef pred_to_angle(pred, epsilon=1e-8):\n    # convert prediction to vector\n    pred_vector = (pred.reshape((-1, bin_num * bin_num, 1)) * angle_bin_vector).sum(axis=1)\n    \n    # normalize\n    pred_vector_norm = np.sqrt((pred_vector**2).sum(axis=1))\n    mask = pred_vector_norm < epsilon\n    pred_vector_norm[mask] = 1\n    \n    # assign <1, 0, 0> to very small vectors (badly predicted)\n    pred_vector /= pred_vector_norm.reshape((-1, 1))\n    pred_vector[mask] = np.array([1., 0., 0.])\n    \n    # convert to angle\n    azimuth = np.arctan2(pred_vector[:, 1], pred_vector[:, 0])\n    azimuth[azimuth < 0] += 2 * np.pi\n    zenith = np.arccos(pred_vector[:, 2])\n    \n    return azimuth, zenith\n\ndef y_to_angle_code(batch_y):\n    azimuth_code = (batch_y[:, 0] > azimuth_edges[1:].reshape((-1, 1))).sum(axis=0)\n    zenith_code = (batch_y[:, 1] > zenith_edges[1:].reshape((-1, 1))).sum(axis=0)\n    angle_code = bin_num * azimuth_code + zenith_code\n    \n    return angle_code","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Data Loading","metadata":{}},{"cell_type":"code","source":"def normalize_data(data):\n    data[:, :, 0] /= 1000   # time\n    data[:, :, 1] /= 300    # charge\n    data[:, :, 3:] /= 600   # space\n    \n    return data\n\ndef prep_validation_data(validation_files_amount):\n    print(\"Processing Validation Data...\")\n\n    # Prepare fixed Validation Set\n    val_x = None\n    val_y = None\n    \n    # Summary\n    print(train_batch_ids[:validation_files_amount])\n\n    # Loop\n    for batch_id in tqdm(train_batch_ids[:validation_files_amount]):\n        val_data_file = np.load(file_format.format(batch_id = batch_id))\n\n        if val_x is None:\n            val_x = val_data_file[\"x\"][:, :, [0,1,2,3,4,5]]\n            val_y = val_data_file[\"y\"]\n        else:\n            val_x = np.append(val_x, val_data_file[\"x\"][:, :, [0,1,2,3,4,5]], axis = 0)\n            val_y = np.append(val_y, val_data_file[\"y\"], axis = 0)\n\n        val_data_file.close()\n        del val_data_file\n        _ = gc.collect()\n\n    # Normalize Data\n    val_x = normalize_data(val_x)\n\n    # Shape Summary\n    print(val_x.shape)\n    \n    return val_x, val_y\n\ndef prep_training_data(start_batch):\n    print(\"Processing Training Data...\")\n    \n    # Placeholders\n    train_x = None\n    train_y = None\n    \n    # Summary\n    train_ids = random.sample(train_batch_ids[start_batch:], train_files_delta)\n    print(train_ids)\n    \n    # Loop\n    for batch_id in tqdm(train_ids):\n        train_data_file = np.load(file_format.format(batch_id = batch_id))\n\n        if train_x is None:\n            train_x = train_data_file[\"x\"][:, :, [0,1,2,3,4,5]]\n            train_y = train_data_file[\"y\"]\n        else:\n            train_x = np.append(train_x, train_data_file[\"x\"][:, :, [0,1,2,3,4,5]], axis = 0)\n            train_y = np.append(train_y, train_data_file[\"y\"], axis = 0)\n\n        train_data_file.close()\n        del train_data_file\n        _ = gc.collect()\n\n    # Normalize data\n    train_x = normalize_data(train_x)\n    \n    # Shape Summary\n    print(train_x.shape)\n    \n    # Output Encoding\n    trn_y_anglecode = y_to_angle_code(train_y)\n        \n    return train_x, trn_y_anglecode","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model","metadata":{}},{"cell_type":"code","source":"def create_model():\n    with strategy.scope(): \n        inputs = tf.keras.layers.Input((pulse_count, feature_count))\n        \n        x = tf.keras.layers.Masking(mask_value = 0., input_shape = (pulse_count, feature_count))(inputs)\n        x = tf.keras.layers.Bidirectional(tf.keras.layers.GRU(lstm_units, return_sequences = True))(x)\n        x = tf.keras.layers.Bidirectional(tf.keras.layers.GRU(lstm_units, return_sequences = True))(x)\n        x = tf.keras.layers.Bidirectional(tf.keras.layers.GRU(lstm_units))(x)        \n        x = tf.keras.layers.Dense(256, activation = 'relu')(x)\n        \n        outputs = tf.keras.layers.Dense(bin_num**2, activation = 'softmax')(x)\n\n        # Finalize Model\n        model = tf.keras.models.Model(inputs = inputs, outputs = outputs)\n\n        # Compile model\n        model.compile(loss = 'sparse_categorical_crossentropy',\n                      optimizer= tf.keras.optimizers.Adam(learning_rate = learning_rate),\n                      metrics = ['accuracy'])\n        \n        # Show Model Summary\n        model.summary()\n\n        return model","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Train Model","metadata":{}},{"cell_type":"code","source":"# Get Fixed Validation Dataset\nval_x, val_y = prep_validation_data(validation_files_amount)\n\n# Create Model\nmodel = create_model()\n\n# For training other than Kaggle environment...provided enough RAM...Load all data\nif data_new_load_interval is None and train_files_delta is None:\n    print('\\nLoading All Train Data')\n    start_batch = validation_files_amount\n    end_batch = start_batch + (len(train_batch_ids) - validation_files_amount)\n    trn_x, trn_y_anglecode = prep_training_data(start_batch, end_batch)\n\n# Epoch Loop\nfor e in range(epochs):\n    print(f'=========== EPOCH: {e}')\n    \n    # Load new random batch of training files .. delta wise .. on Kaggle or Colab with limited RAM.\n    if data_new_load_interval is not None and train_files_delta is not None and e % data_new_load_interval == 0:\n        print(f'\\nLoading Train Data at epoch: {e}')\n        trn_x, trn_y_anglecode = prep_training_data(validation_files_amount)\n    \n    # Number of batches\n    batch_count = trn_x.shape[0] // batch_size\n\n    # Random Shuffle each epoch\n    indices = np.arange(trn_x.shape[0])\n    np.random.shuffle(indices)\n    trn_x = trn_x[indices]\n    trn_y_anglecode = trn_y_anglecode[indices]\n        \n    # Placeholder\n    losses = []\n    accuracy = []\n        \n    # Batch Loop\n    for batch_index in tqdm(range(batch_count), total = batch_count):\n        b_train_x = trn_x[batch_index * batch_size: batch_index * batch_size + batch_size,:]\n        b_train_y = trn_y_anglecode[batch_index * batch_size: batch_index * batch_size + batch_size]\n        \n        metrics = model.train_on_batch(b_train_x, b_train_y)\n        losses.append(metrics[0])\n        accuracy.append(metrics[1])  \n    \n    # Save Model\n    model.save(f'tpu_pp96_n{feature_count}_bin{bin_num}_batch{batch_size}_epoch{e}.h5')\n\n    # Metrics\n    valid_pred = model.predict(val_x, batch_size = batch_size, verbose = verbose)    \n    valid_pred_azimuth, valid_pred_zenith = pred_to_angle(valid_pred)\n    mae = angular_dist_score(val_y[:, 0], val_y[:, 1], valid_pred_azimuth, valid_pred_zenith)    \n    print(f'Total Train Loss: {np.mean(losses):.4f}   Accuracy: {np.mean(accuracy):.4f}  MAE: {mae:.5f}')  \n        \n    # Memory Cleanup\n    gc.collect()","metadata":{"papermill":{"duration":5378.957957,"end_time":"2023-02-27T23:54:04.277619","exception":false,"start_time":"2023-02-27T22:24:25.319662","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]}]}