{"cells":[{"metadata":{},"cell_type":"markdown","source":"## LSTM_decoder_no_map"},{"metadata":{"execution":{"iopub.execute_input":"2020-11-14T09:13:09.327827Z","iopub.status.busy":"2020-11-14T09:13:09.326926Z","iopub.status.idle":"2020-11-14T09:13:13.237584Z","shell.execute_reply":"2020-11-14T09:13:13.236165Z"},"papermill":{"duration":3.948053,"end_time":"2020-11-14T09:13:13.237728","exception":false,"start_time":"2020-11-14T09:13:09.289675","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"# common imports\nimport os\nimport math\nimport time\nimport random\nimport warnings\nwarnings.filterwarnings(\"ignore\")\nimport numpy as np\nimport pandas as pd\nfrom typing import Dict\nfrom tqdm import tqdm\nfrom pathlib import Path\nfrom tempfile import gettempdir\n\n# interactive plot libraries\nimport matplotlib.pyplot as plt\nfrom plotly.offline import init_notebook_mode, iplot # download_plotlyjs, plot\nimport plotly.graph_objs as go\nfrom plotly.subplots import make_subplots\ninit_notebook_mode(connected=True)\n\n# 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.nn.functional as F\n\n#from torch_lr_finder import LRFinder\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","execution_count":null,"outputs":[]},{"metadata":{"execution":{"iopub.execute_input":"2020-11-14T09:13:13.290127Z","iopub.status.busy":"2020-11-14T09:13:13.289336Z","iopub.status.idle":"2020-11-14T09:13:14.093767Z","shell.execute_reply":"2020-11-14T09:13:14.092269Z"},"papermill":{"duration":0.832324,"end_time":"2020-11-14T09:13:14.093953","exception":false,"start_time":"2020-11-14T09:13:13.261629","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"#!rm -rf /kaggle/working/*","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(torch.__version__)","execution_count":null,"outputs":[]},{"metadata":{"execution":{"iopub.execute_input":"2020-11-14T09:13:14.172595Z","iopub.status.busy":"2020-11-14T09:13:14.17015Z","iopub.status.idle":"2020-11-14T09:13:14.178428Z","shell.execute_reply":"2020-11-14T09:13:14.17912Z"},"papermill":{"duration":0.051339,"end_time":"2020-11-14T09:13:14.179295","exception":false,"start_time":"2020-11-14T09:13:14.127956","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"print(l5kit.__version__)","execution_count":null,"outputs":[]},{"metadata":{"execution":{"iopub.execute_input":"2020-11-14T09:13:14.26322Z","iopub.status.busy":"2020-11-14T09:13:14.26253Z","iopub.status.idle":"2020-11-14T09:13:14.270005Z","shell.execute_reply":"2020-11-14T09:13:14.273176Z"},"papermill":{"duration":0.05896,"end_time":"2020-11-14T09:13:14.273389","exception":false,"start_time":"2020-11-14T09:13:14.214429","status":"completed"},"tags":[],"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":{"execution":{"iopub.execute_input":"2020-11-14T09:13:14.357323Z","iopub.status.busy":"2020-11-14T09:13:14.35642Z","iopub.status.idle":"2020-11-14T09:13:14.440934Z","shell.execute_reply":"2020-11-14T09:13:14.441633Z"},"papermill":{"duration":0.131663,"end_time":"2020-11-14T09:13:14.441846","exception":false,"start_time":"2020-11-14T09:13:14.310183","status":"completed"},"tags":[],"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":{"papermill":{"duration":0.039684,"end_time":"2020-11-14T09:13:14.513227","exception":false,"start_time":"2020-11-14T09:13:14.473543","status":"completed"},"tags":[]},"cell_type":"markdown","source":"## Configs"},{"metadata":{"execution":{"iopub.execute_input":"2020-11-14T09:13:14.595175Z","iopub.status.busy":"2020-11-14T09:13:14.594272Z","iopub.status.idle":"2020-11-14T09:13:14.601763Z","shell.execute_reply":"2020-11-14T09:13:14.602874Z"},"papermill":{"duration":0.048864,"end_time":"2020-11-14T09:13:14.603092","exception":false,"start_time":"2020-11-14T09:13:14.554228","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"# --- Lyft configs ---\ncfg = {\n    'format_version': 4,\n    'data_path': \"../input/lyft-motion-prediction-autonomous-vehicles/\",\n    'model_params': {\n        'model_architecture': '',\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_decoder_no_map\",\n        'lr': 5e-4,\n        'weight_path': \"../input/lyftlstm-decoder-no-maps/LSTM_decoder_no_map_1024k.pth\",\n        'lr_find' : False, \n        'train': True,\n        'validate': False,\n        'test': False\n    },\n\n    'raster_params': {\n        'raster_size': [330, 180],\n        'pixel_size': [0.4, 0.4],\n        'ego_center': [0.25, 0.5],\n        'map_type': 'stub_debug',\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': 32,\n        'shuffle': True,\n        'num_workers': 4\n    },\n    \n    'val_data_loader': {\n        'key': 'scenes/validate.zarr',\n        'batch_size': 16,\n        'shuffle': True,\n        'num_workers': 4\n    },\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    'train_params': {\n        'train_start_batch_index' : 32001,\n        'max_num_steps': 32002,\n        'checkpoint_every_n_steps': 2000,\n        'reduction_factor' : 0.9,\n        'step_size' : 1e5\n    }\n}","execution_count":null,"outputs":[]},{"metadata":{"execution":{"iopub.execute_input":"2020-11-14T09:13:14.673878Z","iopub.status.busy":"2020-11-14T09:13:14.673067Z","iopub.status.idle":"2020-11-14T09:13:14.677841Z","shell.execute_reply":"2020-11-14T09:13:14.67849Z"},"papermill":{"duration":0.044685,"end_time":"2020-11-14T09:13:14.678652","exception":false,"start_time":"2020-11-14T09:13:14.633967","status":"completed"},"tags":[],"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_BATCH_SIZE = cfg['train_data_loader']['batch_size'] \nTRAIN_START_BATCH_INDICES = cfg['train_params']['train_start_batch_index']\nEXTENT_RANGE = 5.0 \nMIN_FRAMES_FUTURE = 10","execution_count":null,"outputs":[]},{"metadata":{"execution":{"iopub.execute_input":"2020-11-14T09:13:14.749681Z","iopub.status.busy":"2020-11-14T09:13:14.748721Z","iopub.status.idle":"2020-11-14T09:13:14.751202Z","shell.execute_reply":"2020-11-14T09:13:14.750487Z"},"papermill":{"duration":0.041572,"end_time":"2020-11-14T09:13:14.751327","exception":false,"start_time":"2020-11-14T09:13:14.709755","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"EXTENT_RANGE = 5.0 \nMAX_VELOCITY = 20.0\nMAX_ACCELERATION = 2.0\nMAX_YAW_RATE = np.deg2rad(45)\ndt =cfg['model_params']['history_delta_time']","execution_count":null,"outputs":[]},{"metadata":{"papermill":{"duration":0.0326,"end_time":"2020-11-14T09:13:14.816791","exception":false,"start_time":"2020-11-14T09:13:14.784191","status":"completed"},"tags":[]},"cell_type":"markdown","source":"## Rasterize and initialise"},{"metadata":{"execution":{"iopub.execute_input":"2020-11-14T09:13:14.911391Z","iopub.status.busy":"2020-11-14T09:13:14.909517Z","iopub.status.idle":"2020-11-14T09:13:21.609157Z","shell.execute_reply":"2020-11-14T09:13:21.610686Z"},"papermill":{"duration":6.753796,"end_time":"2020-11-14T09:13:21.610923","exception":false,"start_time":"2020-11-14T09:13:14.857127","status":"completed"},"tags":[],"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":{"papermill":{"duration":0.026594,"end_time":"2020-11-14T09:14:10.767476","exception":false,"start_time":"2020-11-14T09:14:10.740882","status":"completed"},"tags":[]},"cell_type":"markdown","source":"## Loss function"},{"metadata":{"execution":{"iopub.execute_input":"2020-11-14T09:14:10.83846Z","iopub.status.busy":"2020-11-14T09:14:10.837212Z","iopub.status.idle":"2020-11-14T09:14:10.84778Z","shell.execute_reply":"2020-11-14T09:14:10.847083Z"},"papermill":{"duration":0.052976,"end_time":"2020-11-14T09:14:10.847893","exception":false,"start_time":"2020-11-14T09:14:10.794917","status":"completed"},"tags":[],"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":{"papermill":{"duration":0.026633,"end_time":"2020-11-14T09:14:10.90157","exception":false,"start_time":"2020-11-14T09:14:10.874937","status":"completed"},"tags":[]},"cell_type":"markdown","source":"## Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"class LSTM_Decoder(nn.Module):\n    def __init__(self, cfg: Dict, num_modes = 3):\n        super(LSTM_Decoder, self).__init__()\n        self.num_history_frames = cfg[\"model_params\"][\"history_num_frames\"] + 1\n        self.num_future_frames = cfg[\"model_params\"][\"future_num_frames\"]\n        self.num_modes = num_modes\n        \n        # preprocessing layers for inputs\n        self.his_pos_fc   = nn.Linear(2 * self.num_history_frames, 32)\n        self.his_avail_fc = nn.Linear(self.num_history_frames, 16)\n        self.extend_fc    = nn.Linear(3, 3)\n        \n        # LSTM layers   \n        self.embed_size  = 3 * (self.num_modes) # 2 positions and 1 confidence for each mode\n        self.hidden_size = 32 + 16 + 3\n        self.decoder     = nn.LSTMCell(self.embed_size, self.hidden_size, bias=True)\n        self.lstm_tgt_fc = nn.Linear(self.hidden_size, self.embed_size)\n        \n    def prepare_data(self, data, device):\n        target_availabilities = data[\"target_availabilities\"].to(device)\n        targets = data[\"target_positions\"].to(device)\n        bs = targets.shape[0]\n\n        extent = (data['extent'].to(device) / EXTENT_RANGE).view(bs, -1)\n        history_positions = data['history_positions'].to(device).view(bs, -1)\n        history_availabilities = data['history_availabilities'].to(device).view(bs, -1)\n        return [bs, \n                history_positions, history_availabilities, extent,\n                targets, target_availabilities]\n    \n    def LSTM_decoder_step(self, bs, his_pos, his_avail, bb, tgts, tgt_avail, criterion):\n        # preallocating memory for output\n        LSTM_output = torch.zeros(bs, self.num_future_frames, self.embed_size).to(device)\n        # first decoder input\n        initial_input = torch.zeros(bs, self.embed_size).to(device) \n        loss = 0.0\n        \n        # initial hidden state\n        in1 = model.his_pos_fc(his_pos)\n        in2 = model.his_avail_fc(his_avail)\n        in3 = model.extend_fc(bb)\n        decoder_hidden = torch.cat((in1, in2, in3), dim=1)\n        decoder_cellState = torch.zeros_like(decoder_hidden)\n        \n        # Decode hidden state in future trajectory\n        decoder_input = initial_input        \n        for t in range(self.num_future_frames):\n            decoder_hidden, decoder_cellState = self.decoder(decoder_input, (decoder_hidden, decoder_cellState))\n            decoder_output = self.lstm_tgt_fc(decoder_hidden)\n            LSTM_output[:, t, :] = decoder_output\n\n            # Use own predictions as inputs at next step\n            decoder_input = decoder_output\n\n        pred, confidences = torch.split(LSTM_output, 2 * self.num_modes, dim=2)\n        pred = pred.reshape(bs, self.num_modes, self.num_future_frames, 2)\n        confidences = torch.softmax(torch.mean(confidences, dim =1), dim=1)\n        loss = criterion(tgts, pred, confidences,tgt_avail)\n        return loss, pred, confidences\n    \n    # create masks for history data \n    def mask_history_data(self, device, history_positions, history_availabilities):\n        # mask for history availabbilities and history_positions for reducing overfit\n        avail_mask = torch.rand(history_availabilities.shape).to(device) > 0.9\n        pos_mask = torch.repeat_interleave(avail_mask, 2, dim=1).to(device)\n        history_availabilities = history_availabilities.masked_fill(avail_mask, 0)\n        history_positions = history_positions.masked_fill(pos_mask, 0)\n        return history_positions, history_availabilities\n    \n    # inference / prediction pipeline\n    def forward(self, data, device, criterion = pytorch_neg_multi_log_likelihood_batch):\n        bs, his_pos, his_avail, bb, tgts, tgt_avail = self.prepare_data(data, device)\n        return self.LSTM_decoder_step(bs, his_pos, his_avail, bb, tgts, tgt_avail, criterion)\n    \n    # Training forward function with addtional masks for history position and availability data\n    def train_forward(self, data, device, criterion = pytorch_neg_multi_log_likelihood_batch):\n        bs, his_pos, his_avail, bb, tgts, tgt_avail = self.prepare_data(data, device)\n        his_pos, his_avail = self.mask_history_data(device, his_pos, his_avail)\n        #print(his_pos[0])\n        return self.LSTM_decoder_step(bs, his_pos, his_avail, bb, tgts, tgt_avail, criterion)","execution_count":null,"outputs":[]},{"metadata":{"execution":{"iopub.execute_input":"2020-11-14T09:14:11.48009Z","iopub.status.busy":"2020-11-14T09:14:11.479175Z","iopub.status.idle":"2020-11-14T09:14:11.483613Z","shell.execute_reply":"2020-11-14T09:14:11.484126Z"},"papermill":{"duration":0.411044,"end_time":"2020-11-14T09:14:11.484278","exception":false,"start_time":"2020-11-14T09:14:11.073234","status":"completed"},"tags":[],"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":"model = LSTM_Decoder(cfg)\nmodel.to(device)\noptimizer = optim.Adam(model.parameters(), lr=cfg[\"model_params\"][\"lr\"])\ncriterion = pytorch_neg_multi_log_likelihood_batch\nprint(f'device {device}')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#load weight if there is a pretrained model\nweight_path = cfg[\"model_params\"][\"weight_path\"]\nif weight_path != \"\":\n    checkpoint = torch.load(weight_path)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#if weight_path != \"\":\nmodel.load_state_dict(checkpoint['state_dict'])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#if weight_path != \"\":\noptimizer.load_state_dict(checkpoint['optimizer'])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(find_no_of_trainable_params(model))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print('Before lr : ' , optimizer.param_groups[0]['lr'])\noptimizer.param_groups[0]['lr'] = cfg[\"model_params\"][\"lr\"]\nprint('After lr  : ' , optimizer.param_groups[0]['lr'])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"scheduler = optim.lr_scheduler.OneCycleLR(optimizer, max_lr= cfg[\"model_params\"][\"lr\"],\n                                          total_steps=cfg['train_params']['max_num_steps'],\n                                          pct_start=0.4, div_factor=10)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Train Dataloader"},{"metadata":{"execution":{"iopub.execute_input":"2020-11-14T09:13:21.670706Z","iopub.status.busy":"2020-11-14T09:13:21.669796Z","iopub.status.idle":"2020-11-14T09:14:09.446071Z","shell.execute_reply":"2020-11-14T09:14:09.445115Z"},"papermill":{"duration":47.805745,"end_time":"2020-11-14T09:14:09.446194","exception":false,"start_time":"2020-11-14T09:13:21.640449","status":"completed"},"tags":[],"trusted":true},"cell_type":"code","source":"# ===== INIT TRAIN DATASET============================================================\nif (cfg['model_params']['train'] == True) or (cfg['model_params']['lr_find'] == True):\n    train_cfg = cfg[\"train_data_loader\"]\n    train_zarr = ChunkedDataset(dm.require(train_cfg[\"key\"])).open()\n    train_dataset = AgentDataset(cfg, train_zarr, rasterizer, \n                                 min_frame_history = cfg['model_params']['history_num_frames'],\n                                 min_frame_future=MIN_FRAMES_FUTURE)\n    \n    print('Length of Train dataset is ' ,len(train_dataset))\n    #print(\"==================================TRAIN DATA==================================\")\n    #print(train_dataset)\n    \n    sampled_indices = np.random.choice(len(train_dataset), size = len(train_dataset), replace = False)\n    #print('Before slicing, start indices are ', sampled_indices[0:5])\n    #print('TRAIN_START_INDICES', TRAIN_START_BATCH_INDICES * TRAIN_BATCH_SIZE)\n    \n    sampled_indices = sampled_indices[TRAIN_START_BATCH_INDICES * TRAIN_BATCH_SIZE:]\n    #print('After slicing, start indices are ', sampled_indices[0:5])\n    \n    Datasampler = SubsetRandomSampler(sampled_indices)\n    train_dataloader = DataLoader(train_dataset, sampler=Datasampler, batch_size=train_cfg[\"batch_size\"], \n                             num_workers=train_cfg[\"num_workers\"])\n    print('Length of train_dataloader is ' ,len(train_dataloader))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Validation Dataloader"},{"metadata":{"trusted":true},"cell_type":"code","source":"eval_base_path = '../input/lyft-validation-chopped-100/validate_chopped_100'\neval_cfg = cfg[\"val_data_loader\"]\neval_zarr_path = str(Path(eval_base_path) / Path(dm.require(eval_cfg[\"key\"])).name)\neval_mask_path = str(Path(eval_base_path) / \"mask.npz\")\neval_gt_path = str(Path(eval_base_path) / \"gt.csv\")\n\neval_zarr = ChunkedDataset(eval_zarr_path).open()\neval_mask = np.load(eval_mask_path)[\"arr_0\"]\n# ===== INIT DATASET AND LOAD MASK\neval_dataset = AgentDataset(cfg, eval_zarr, rasterizer, agents_mask=eval_mask)\neval_dataloader = DataLoader(eval_dataset, shuffle=eval_cfg[\"shuffle\"], batch_size=eval_cfg[\"batch_size\"], \n                             num_workers=eval_cfg[\"num_workers\"])\nprint(eval_dataset)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def model_validation_score(model, pred_path):\n    # ==== EVAL LOOP\n    model.eval()\n    torch.set_grad_enabled(False)\n\n    # store information for evaluation\n    future_coords_offsets_pd = []\n    timestamps = []\n    confidences_list = []\n    agent_ids = []\n    progress_bar = tqdm(eval_dataloader)\n\n    for data in progress_bar:\n\n        _, preds, confidences = model.forward(data, device)\n\n        #fix for the new environment\n        preds = preds.cpu().numpy()\n        world_from_agents = data[\"world_from_agent\"].numpy()\n        centroids = data[\"centroid\"].numpy()\n        coords_offset = []\n\n        # convert into world coordinates and compute offsets\n        for idx in range(len(preds)):\n            for mode in range(3):\n                preds[idx, mode, :, :] = transform_points(preds[idx, mode, :, :], world_from_agents[idx]) - centroids[idx][:2]\n\n        future_coords_offsets_pd.append(preds.copy())\n        confidences_list.append(confidences.cpu().numpy().copy())\n        timestamps.append(data[\"timestamp\"].numpy().copy())\n        agent_ids.append(data[\"track_id\"].numpy().copy())  \n    \n    write_pred_csv(pred_path,\n               timestamps=np.concatenate(timestamps),\n               track_ids=np.concatenate(agent_ids),\n               coords=np.concatenate(future_coords_offsets_pd),\n               confs=np.concatenate(confidences_list),\n              )\n    \n    metrics = compute_metrics_csv(eval_gt_path, pred_path, [neg_multi_log_likelihood, time_displace])\n    for metric_name, metric_mean in metrics.items():\n        print(metric_name, metric_mean)\n    \n    #return [future_coords_offsets_pd, confidences_list, timestamps, agent_ids]","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## LR_finder "},{"metadata":{"trusted":true},"cell_type":"code","source":"def plot_lr_finder_results(lr_finder): \n    # Create subplot grid\n    fig = make_subplots(rows=1, cols=2)\n    # layout ={'title': 'Lr_finder_result'}\n    \n    # Create a line (trace) for the lr vs loss, gradient of loss\n    trace0 = go.Scatter(x=lr_finder['log_lr'], y=lr_finder['smooth_loss'],name='log_lr vs smooth_loss')\n    trace1 = go.Scatter(x=lr_finder['log_lr'], y=lr_finder['grad_loss'],name='log_lr vs loss gradient')\n\n    # Add subplot trace & assign to each grid\n    fig.add_trace(trace0, row=1, col=1)\n    fig.add_trace(trace1, row=1, col=2)\n    iplot(fig, show_link=False)\n    #fig.write_html(cfg['model_params']['weight_path'] + '_lr_find.html')\n\ndef find_lr(data_loader, init_value = 1e-8, final_value=1.0, beta = 0.98, num_batches = 200):\n    assert(num_batches > 0)\n    mult = (final_value / init_value) ** (1/num_batches)\n    lr = init_value\n    optimizer.param_groups[0]['lr'] = lr\n    batch_num = 0\n    avg_loss = 0.0\n    best_loss = 0.0\n    smooth_losses = []\n    raw_losses = []\n    log_lrs = []\n    dataloader_it = iter(data_loader)\n    progress_bar = tqdm(range(num_batches))\n    \n    for idx in progress_bar:\n        batch_num += 1\n        try:\n            data = next(dataloader_it)\n        except StopIteration:\n            dataloader_it = iter(data_loader)\n            data = next(dataloader_it)\n        \n        # Forward pass\n        model.train()\n        torch.set_grad_enabled(True)\n        \n        # handle exception in criterion\n        try:\n            loss, _, _ = model.train_forward(data, device)\n        except:\n            if len(smooth_losses) > 1:\n                grad_loss = np.gradient(smooth_losses)\n            else:\n                grad_loss = 0.0\n            lr_finder_results = {'log_lr':log_lrs, 'raw_loss':raw_losses, \n                                 'smooth_loss':smooth_losses, 'grad_loss': grad_loss}\n            return lr_finder_results\n        \n        #Compute the smoothed loss\n        avg_loss = beta * avg_loss + (1-beta) *loss.item()\n        smoothed_loss = avg_loss / (1 - beta**batch_num)\n        \n        #Stop if the loss is exploding\n        if batch_num > 1 and smoothed_loss > 50 * best_loss:\n            if len(smooth_losses) > 1:\n                grad_loss = np.gradient(smooth_losses)\n            else:\n                grad_loss = 0.0\n            lr_finder_results = {'log_lr':log_lrs, 'raw_loss':raw_losses, \n                                 'smooth_loss':smooth_losses, 'grad_loss': grad_loss}\n            return lr_finder_results\n        \n        #Record the best loss\n        if smoothed_loss < best_loss or batch_num==1:\n            best_loss = smoothed_loss\n        \n        #Store the values\n        raw_losses.append(loss.item())\n        smooth_losses.append(smoothed_loss)\n        log_lrs.append(math.log10(lr))\n        \n        # Backward pass\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n        \n        # print info\n        progress_bar.set_description(f\"loss: {loss.item()},smoothed_loss: {smoothed_loss},lr : {lr}\")\n\n        #Update the lr for the next step\n        lr *= mult\n        optimizer.param_groups[0]['lr'] = lr\n    \n    grad_loss = np.gradient(lr_finder_results['smooth_loss'])\n    lr_finder_results = {'log_lr':log_lrs, 'raw_loss':raw_losses, \n                         'smooth_loss':smooth_losses, '': grad_loss}\n    return lr_finder_results","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if cfg['model_params']['lr_find'] == True:\n    lr_finder_results = find_lr(train_dataloader)\n    plot_lr_finder_results(lr_finder_results)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Training loop"},{"metadata":{"trusted":true},"cell_type":"code","source":"# ==== TRAINING LOOP =========================================================\nif cfg[\"model_params\"][\"train\"] == True:\n    \n    print('TRAINING ABOUT TO START ... FROM ', TRAIN_START_BATCH_INDICES, \n      '  BATCH AND FOR ', cfg['train_params']['max_num_steps'], ' BATCHES', ' WITH BATCH SIZE', TRAIN_BATCH_SIZE)\n    \n    tr_it = iter(train_dataloader)\n    progress_bar = tqdm(range(TRAIN_START_BATCH_INDICES, \n                              TRAIN_START_BATCH_INDICES + cfg[\"train_params\"][\"max_num_steps\"]))\n    num_iter = cfg[\"train_params\"][\"max_num_steps\"]\n    losses_train = []\n    smooth_losses = []\n    iterations = []\n    metrics = []\n    times = []\n    lr_list = []\n    model_name = cfg[\"model_params\"][\"model_name\"]\n    start = time.time()\n    iteration = 0\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, _, _ = model.train_forward(data, device)\n\n        # Backward pass\n        optimizer.zero_grad()\n        loss.backward()\n        \n        #Gradient Value Clipping\n        nn.utils.clip_grad_norm_(model.parameters(),max_norm=5.0)\n        optimizer.step()\n        \n        # lr scheduler\n        scheduler.step()\n        \n        losses_train.append(loss.item())\n        lr_list.append(optimizer.param_groups[0]['lr'])\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            sample_number = i * cfg['train_data_loader']['batch_size'] \n            \n            #pred_path = f\"{gettempdir()}/pred.csv\"\n            #model_validation_score(model, pred_path)\n            \n            state = {\n              'state_dict': model.state_dict(),\n              'optimizer': optimizer.state_dict()\n            }\n            torch.save(state, f'{model_name}_{sample_number}k.pth')\n            iterations.append(i)\n            metrics.append(np.mean(losses_train))\n            times.append((time.time()-start)/60)\n\n    sample_number = i * cfg['train_data_loader']['batch_size']\n    results = pd.DataFrame({'iterations': iterations, 'metrics (avg)': metrics, 'elapsed_time (mins)': times})\n    results.to_csv(f\"train_metrics_{model_name}_{sample_number}k.csv\", index = False)\n    train_losses_csv = pd.DataFrame({'iteration': TRAIN_START_BATCH_INDICES * TRAIN_BATCH_SIZE + np.arange(len(losses_train)), \n                                 'losses_train': losses_train,\n                                 'lr_list' : lr_list})\n    train_losses_csv.to_csv(f\"train_losses_{model_name}_{sample_number}k.csv\", index = False)\n    print(f\"Total training time is {(time.time()-start)/60} mins\")\n    print(results.head())","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Validation loop"},{"metadata":{"trusted":true},"cell_type":"code","source":"if cfg[\"model_params\"][\"validate\"] == True:\n    pred_path = f\"{gettempdir()}/pred.csv\"\n    model_validation_score(model, pred_path)","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}