{"cells":[{"metadata":{},"cell_type":"markdown","source":"## LSTM method for Motion prediction"},{"metadata":{"trusted":true},"cell_type":"markdown","source":"Objective of the competition is to predict the future trajectories of other vehicles using the past information (Bird's eye view of the scene containing agents detected by perception system, past trajectories lane information, traffic lights etc).\n\n\nSince this is sequential problem, I thought of using LSTM based models. Basic idea is as follows:\n\n[LSTM_baseline idea](https://www.kaggle.com/suryajrrafl/lstm-baseline-weights?select=lstm+baseline+idea.jpg)\n\nThis is my first attempt at kaggle competition, pytorch and LSTM models. Suggestions are most welcome.\n\n\n**REFERENCES**\n\nSome helper functions in this notebook were taken from the great public kernels avaiable. \n\n[Great reference notebook using Resnet model](https://www.kaggle.com/huanvo/lyft-complete-train-and-prediction-pipeline)\n\n[Pytorch baseline train](https://www.kaggle.com/pestipeti/pytorch-baseline-train)"},{"metadata":{},"cell_type":"markdown","source":"## Import libraries"},{"metadata":{"trusted":true},"cell_type":"code","source":"# torch imports\nimport torch\nfrom torch import nn, optim\nfrom torch.utils.data import DataLoader, SubsetRandomSampler\nfrom torchvision.models.resnet import resnet50, resnet18, resnet34, resnet101\nimport torch.functional as F\n\n# l5kit imports\nimport l5kit\nfrom l5kit.configs import load_config_data\nfrom l5kit.data import LocalDataManager, ChunkedDataset\nfrom l5kit.dataset import AgentDataset, EgoDataset\nfrom l5kit.rasterization import build_rasterizer\nfrom l5kit.evaluation import write_pred_csv, compute_metrics_csv, read_gt_csv, create_chopped_dataset\nfrom l5kit.evaluation.chop_dataset import MIN_FUTURE_STEPS\nfrom l5kit.evaluation.metrics import neg_multi_log_likelihood, time_displace\nfrom l5kit.geometry import transform_points\nfrom l5kit.visualization import PREDICTED_POINTS_COLOR, TARGET_POINTS_COLOR, draw_trajectory\n\n# common imports\nimport os\nimport random\nimport time\nimport pandas as pd\nfrom typing import Dict\nfrom tempfile import gettempdir\nimport matplotlib.pyplot as plt\nimport numpy as np\nfrom tqdm import tqdm\nfrom prettytable import PrettyTable\nfrom pathlib import Path\nimport warnings\nwarnings.filterwarnings(\"ignore\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"l5kit.__version__","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.cuda.is_available()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Helper functions"},{"metadata":{},"cell_type":"markdown","source":"### Loss function"},{"metadata":{"trusted":true},"cell_type":"code","source":"# --- Function utils ---\n# Original code from https://github.com/lyft/l5kit/blob/20ab033c01610d711c3d36e1963ecec86e8b85b6/l5kit/l5kit/evaluation/metrics.py\nfrom torch import Tensor\n\n\ndef pytorch_neg_multi_log_likelihood_batch(\n    gt: Tensor, pred: Tensor, confidences: Tensor, avails: Tensor\n) -> Tensor:\n    \"\"\"\n    Compute a negative log-likelihood for the multi-modal scenario.\n    log-sum-exp trick is used here to avoid underflow and overflow, For more information about it see:\n    https://en.wikipedia.org/wiki/LogSumExp#log-sum-exp_trick_for_log-domain_calculations\n    https://timvieira.github.io/blog/post/2014/02/11/exp-normalize-trick/\n    https://leimao.github.io/blog/LogSumExp/\n    Args:\n        gt (Tensor): array of shape (bs)x(time)x(2D coords)\n        pred (Tensor): array of shape (bs)x(modes)x(time)x(2D coords)\n        confidences (Tensor): array of shape (bs)x(modes) with a confidence for each mode in each sample\n        avails (Tensor): array of shape (bs)x(time) with the availability for each gt timestep\n    Returns:\n        Tensor: negative log-likelihood for this example, a single float number\n    \"\"\"\n    assert len(pred.shape) == 4, f\"expected 3D (MxTxC) 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 2D (Time x Coords) array for gt, got {gt.shape}\"\n    assert confidences.shape == (batch_size, num_modes), f\"expected 1D (Modes) array for gt, got {confidences.shape}\"\n    assert torch.allclose(torch.sum(confidences, dim=1), confidences.new_ones((batch_size,))), \"confidences should sum to 1\"\n    assert avails.shape == (batch_size, future_len), f\"expected 1D (Time) array for gt, got {avails.shape}\"\n    # assert all data are valid\n    assert torch.isfinite(pred).all(), \"invalid value found in pred\"\n    assert torch.isfinite(gt).all(), \"invalid value found in gt\"\n    assert torch.isfinite(confidences).all(), \"invalid value found in confidences\"\n    assert torch.isfinite(avails).all(), \"invalid value found in avails\"\n\n    # convert to (batch_size, num_modes, future_len, num_coords)\n    gt = torch.unsqueeze(gt, 1)  # add modes\n    avails = avails[:, None, :, None]  # add modes and cords\n\n    # error (batch_size, num_modes, future_len)\n    error = torch.sum(((gt - pred) * avails) ** 2, dim=-1)  # reduce coords and use availability\n\n    with np.errstate(divide=\"ignore\"):  # when confidence is 0 log goes to -inf, but we're fine with it\n        # error (batch_size, num_modes)\n        error = torch.log(confidences) - 0.5 * torch.sum(error, dim=-1)  # reduce time\n\n    # use max aggregator on modes for numerical stability\n    # error (batch_size, num_modes)\n    max_value, _ = error.max(dim=1, keepdim=True)  # error are negative at this point, so max() gives the minimum one\n    error = -torch.log(torch.sum(torch.exp(error - max_value), dim=-1, keepdim=True)) - max_value  # reduce modes\n    # print(\"error\", error)\n    return torch.mean(error)\n\n\ndef pytorch_neg_multi_log_likelihood_single(\n    gt: Tensor, pred: Tensor, avails: Tensor\n) -> Tensor:\n    \"\"\"\n\n    Args:\n        gt (Tensor): array of shape (bs)x(time)x(2D coords)\n        pred (Tensor): array of shape (bs)x(time)x(2D coords)\n        avails (Tensor): array of shape (bs)x(time) with the availability for each gt timestep\n    Returns:\n        Tensor: negative log-likelihood for this example, a single float number\n    \"\"\"\n    # pred (bs)x(time)x(2D coords) --> (bs)x(mode=1)x(time)x(2D coords)\n    # create confidence (bs)x(mode=1)\n    batch_size, future_len, num_coords = pred.shape\n    confidences = pred.new_ones((batch_size, 1))\n    return pytorch_neg_multi_log_likelihood_batch(gt, pred.unsqueeze(1), confidences, avails)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Random seed generation function"},{"metadata":{"trusted":true},"cell_type":"code","source":"def set_seed(seed):\n    random.seed(seed)\n    np.random.seed(seed)\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    \nset_seed(42)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Resnet forward function"},{"metadata":{"trusted":true},"cell_type":"code","source":"def resnet_forward(backbone, x):    \n    #with torch.set_grad_enabled(False):\n    with torch.no_grad():\n        x = backbone.conv1(x)\n        x = backbone.bn1(x)\n        x = backbone.relu(x)\n        x = backbone.maxpool(x)\n\n        x = backbone.layer1(x)\n        x = backbone.layer2(x)\n        x = backbone.layer3(x)\n        x = backbone.layer4(x)\n\n        x = backbone.avgpool(x)\n        x = torch.flatten(x, 1)\n    return x","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Number of trainable parameters in model"},{"metadata":{"trusted":true},"cell_type":"code","source":"def find_no_of_trainable_params(model):\n    total_trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n    #print(total_trainable_params)\n    return total_trainable_params","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### LSTM input creation function"},{"metadata":{"trusted":true},"cell_type":"code","source":"def LSTM_batch_transform(image_data, base_model):    \n    \n    BATCH_SIZE = image_data.shape[0]\n    \n    \"\"\" LANES, TRAFFIC LIGHT DATA ENCODING \"\"\"\n    infra_data = image_data[:, -3:, :, :]\n    infra_data = resnet_forward(base_model, infra_data)\n    infra_data = torch.repeat_interleave(infra_data.unsqueeze(1), NUMBER_OF_HISTORY_FRAMES, dim=1)\n    #print(infra_data.shape)\n    \n    \"\"\" EGO, AGENT VEHICLE DATA ENCODING \"\"\"\n    # agent frames\n    agent_data = image_data[:, 0:NUMBER_OF_HISTORY_FRAMES, :, :]\n    #print(agent_data.shape)\n\n    # ego vehicle frames\n    ego_data = image_data[:, NUMBER_OF_HISTORY_FRAMES:-3, :,:]\n    #print(ego_data.shape)\n\n    # combined ego and agent frames, duplicating across 3 channels\n    vehicle_data = torch.repeat_interleave(ego_data + agent_data, 3, dim=1)\n\n    # pretrained model requires (batch_size, 3, 224, 224), hence reshaping\n    vehicle_data = vehicle_data.view(-1, 3, RASTER_IMG_SIZE, RASTER_IMG_SIZE)\n    #print(vehicle_data.shape)\n\n    # passing through model and reshaping\n    history_vehicle_data = resnet_forward(base_model, vehicle_data)\n    history_vehicle_data =  history_vehicle_data.view(BATCH_SIZE, -1, 512)\n    #print(history_vehicle_data.shape)\n    \n    \"\"\"concatenating history_vehicle_data and infra_data \"\"\"\n    LSTM_input = torch.cat((history_vehicle_data, infra_data), dim=-1)\n    #print(f'LSTM input shape is {temp.shape}')\n    \n    return LSTM_input","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Model Forward pass function"},{"metadata":{"trusted":true},"cell_type":"code","source":"def forward(data, model, hidden_state, device, criterion = pytorch_neg_multi_log_likelihood_batch):\n    inputs = data[\"image\"].to(device)\n    target_availabilities = data[\"target_availabilities\"].to(device)\n    targets = data[\"target_positions\"].to(device)\n    batch_size = inputs.shape[0]\n    \n    # converting image data to sequential data for LSTM model\n    LSTM_input = LSTM_batch_transform(inputs, encoding_model)\n    \n    # LSTM model prediction and confidence\n    prediction, hidden_state = model(LSTM_input, hidden_state)\n    hidden_state = (hidden_state[0].data, hidden_state[1].data)\n    prediction, confidences = torch.split(prediction, 300, dim=1)\n    prediction = prediction.view(batch_size, 3, 50, 2)\n    confidences = torch.softmax(confidences, dim=1)\n    \n    # calculating NLL loss \n    loss = pytorch_neg_multi_log_likelihood_batch(targets, prediction, confidences, target_availabilities)\n    \n    return loss, hidden_state, prediction, confidences","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Base LSTM Structure"},{"metadata":{"trusted":true},"cell_type":"code","source":"class RNN(nn.Module):\n    def __init__(self, input_size, output_size, hidden_dim, n_layers, time_steps, use_LSTM = False):\n        super(RNN, self).__init__()\n        \n        self.hidden_dim = hidden_dim\n        self.time_steps = time_steps\n        \n        # define an RNN with specified parameters\n        # batch_first means that the first dim of the input and output will be the batch_size\n        \n        if use_LSTM == True:\n            self.rnn = nn.LSTM(input_size, hidden_dim, n_layers, batch_first=True)\n        else:\n            self.rnn = nn.RNN(input_size, hidden_dim, n_layers, batch_first=True)\n        \n        \n        # last, fully-connected layer\n        self.fc = nn.Linear(time_steps * hidden_dim, output_size)\n\n    def forward(self, x, hidden):\n        # x (batch_size, time_step, input_size)\n        # hidden (n_layers, batch_size, hidden_dim)\n        # r_out (batch_size, time_step, hidden_size)\n        batch_size = x.size(0)\n        \n        # get RNN outputs\n        r_out, hidden = self.rnn(x, hidden)\n        # shape output to be (batch_size*time_step, hidden_dim)\n        r_out = r_out.reshape(batch_size,-1)  \n        \n        # get final output \n        output = self.fc(r_out)\n        \n        return output, hidden","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Configs"},{"metadata":{"trusted":true},"cell_type":"code","source":"# --- Lyft configs ---\ncfg = {\n    'format_version': 4,\n    'data_path': \"/kaggle/input/lyft-motion-prediction-autonomous-vehicles/\",\n    'model_params': {\n        'model_architecture': 'LSTM',\n        'history_num_frames': 10,\n        'history_step_size': 1,\n        'history_delta_time': 0.1,\n        'future_num_frames': 50,\n        'future_step_size': 1,\n        'future_delta_time': 0.1,\n        'model_name': \"LSTM_baseline_r34\",\n        'weight_path': \"/kaggle/input/lstm-baseline-weights/LSTM_baseline_r34_9750.pth\",\n        'lr': 1e-3,\n        'train': True,\n        'predict': False\n    },\n\n    'raster_params': {\n        'raster_size': [224, 224],\n        'pixel_size': [0.5, 0.5],\n        'ego_center': [0.25, 0.5],\n        'map_type': 'py_semantic',\n        'satellite_map_key': 'aerial_map/aerial_map.png',\n        'semantic_map_key': 'semantic_map/semantic_map.pb',\n        'dataset_meta_key': 'meta.json',\n        'filter_agents_threshold': 0.5\n    },\n\n    'train_data_loader': {\n        'key': 'scenes/train.zarr',\n        'batch_size': 16,\n        'shuffle': True,\n        'num_workers': 4\n    },\n    \n    'test_data_loader': {\n        'key': 'scenes/test.zarr',\n        'batch_size': 32,\n        'shuffle': False,\n        'num_workers': 4\n    },\n\n    'sample_data_loader': {\n        'key': 'scenes/sample.zarr',\n        'batch_size': 16,\n        'shuffle': True,\n        'num_workers': 4\n    },\n\n    'train_params': {\n        'train_start_index' : 9751,\n        'max_num_steps': 12002,\n        'checkpoint_every_n_steps': 500,\n    }\n}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Constants used for training"},{"metadata":{"trusted":true},"cell_type":"code","source":"NUMBER_OF_HISTORY_FRAMES = cfg['model_params']['history_num_frames'] + 1\nRASTER_IMG_SIZE = cfg['raster_params']['raster_size'][0]\nNUM_MODES = 3\nNUMBER_OF_FUTURE_FRAMES = cfg['model_params']['future_num_frames']\n\n### TRAIN FROM WHERE LEFT OFF, CHANGE THE STARTING INDICES VARIABLE ACCORDINGLY\nTRAIN_START_INDICES = cfg['train_params']['train_start_index']","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Load the training data"},{"metadata":{"trusted":true},"cell_type":"code","source":"# set env variable for data\nDIR_INPUT = cfg[\"data_path\"]\nos.environ[\"L5KIT_DATA_FOLDER\"] = DIR_INPUT\ndm = LocalDataManager(None)\nrasterizer = build_rasterizer(cfg, dm)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# ===== INIT TRAIN DATASET============================================================\ntrain_cfg = cfg[\"train_data_loader\"]\ntrain_zarr = ChunkedDataset(dm.require(train_cfg[\"key\"])).open()\ntrain_dataset = AgentDataset(cfg, train_zarr, rasterizer)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print('Length of Train dataset is ' ,len(train_dataset))\nprint(\"==================================TRAIN DATA==================================\")\nprint(train_dataset)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"len(train_dataset)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sampled_indices = np.random.choice(len(train_dataset), size = len(train_dataset), replace = False)\nprint('Before slicing, start indices are ', sampled_indices[0:10])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"TRAIN_START_INDICES","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sampled_indices = sampled_indices[TRAIN_START_INDICES:]\nprint('After slicing, start indices are ', sampled_indices[0:10])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"Datasampler = SubsetRandomSampler(sampled_indices)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dataloader = DataLoader(train_dataset, sampler=Datasampler, batch_size=train_cfg[\"batch_size\"], \n                             num_workers=train_cfg[\"num_workers\"])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## CUDA device && encoding model"},{"metadata":{"trusted":true},"cell_type":"code","source":"# ==== INIT MODEL=================\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\nprint(f'device {device}')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"encoding_model = resnet34(pretrained=True)\nencoding_model.to(device);\n\n# Freeze parameters so we don't backprop through them\nfor param in encoding_model.parameters():\n    param.requires_grad = False\n\nTotal_trainable_params = find_no_of_trainable_params(encoding_model)\nprint(f'There are {Total_trainable_params} trainable parameters in the model')\n\n# set to evaluation mode\nencoding_model.eval();","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## LSTM model, Optimiser, criterion"},{"metadata":{"trusted":true},"cell_type":"code","source":"# decide on hyperparameters\ninput_size   = 1024 \noutput_size  = 303\nhidden_dim   = 64\nn_layers     = 2","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# instantiate an RNN\nmodel = RNN(input_size, output_size, hidden_dim, n_layers, 11, use_LSTM=True)\nmodel.to(device)\n#print(LSTM_baseline_model)\n\ntotal_params = find_no_of_trainable_params(model)\nprint(f'There are {total_params} parameters in the LSTM model')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"## loading the pretrained weights\nmodel.load_state_dict(torch.load(cfg['model_params']['weight_path']))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"## Adam optimiser function\noptimizer = optim.Adam(model.parameters(), lr=cfg[\"model_params\"][\"lr\"])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"# ==== TRAINING LOOP =========================================================\nif cfg[\"model_params\"][\"train\"]:\n    \n    tr_it = iter(train_dataloader)\n    progress_bar = tqdm(range(TRAIN_START_INDICES, \n                              TRAIN_START_INDICES + cfg[\"train_params\"][\"max_num_steps\"]))\n    num_iter = cfg[\"train_params\"][\"max_num_steps\"]\n    losses_train = []\n    iterations = []\n    metrics = []\n    times = []\n    model_name = cfg[\"model_params\"][\"model_name\"]\n    start = time.time()\n    hidden_state = None\n    \n    for i in progress_bar:\n        try:\n            data = next(tr_it)\n        except StopIteration:\n            tr_it = iter(train_dataloader)\n            data = next(tr_it)\n            \n        # Forward pass\n        model.train()\n        torch.set_grad_enabled(True)\n        loss, hidden_state, _, _ = forward(data, model, hidden_state, device)\n\n        # Backward pass\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n\n        losses_train.append(loss.item())\n\n        progress_bar.set_description(f\"loss: {loss.item()} loss(avg): {np.mean(losses_train)}\")\n        if i % cfg['train_params']['checkpoint_every_n_steps'] == 0:\n            torch.save(model.state_dict(), f'{model_name}_{i + TRAIN_START_INDICES}.pth')\n            iterations.append(i)\n            metrics.append(np.mean(losses_train))\n            times.append((time.time()-start)/60)\n    \n    results = pd.DataFrame({'iterations': iterations, 'metrics (avg)': metrics, 'elapsed_time (mins)': times})\n    results.to_csv(f\"train_metrics_{model_name}_{num_iter}.csv\", index = False)\n    train_losses_csv = pd.DataFrame({'iteration': TRAIN_START_INDICES + np.arange(len(losses_train)), \n                                 'losses_train': losses_train})\n    train_losses_csv.to_csv(f\"train_losses_{model_name}_{num_iter}.csv\", index = False)\n    print(f\"Total training time is {(time.time()-start)/60} mins\")\n    print(results.head())","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}