{"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":"code","source":"import sys\nsys.path.insert(0, '/kaggle/input/kbl5kit-v2/l5kit/')","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2021-06-08T07:48:59.002518Z","iopub.execute_input":"2021-06-08T07:48:59.003106Z","iopub.status.idle":"2021-06-08T07:48:59.008244Z","shell.execute_reply.started":"2021-06-08T07:48:59.003057Z","shell.execute_reply":"2021-06-08T07:48:59.007245Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# from IPython.core.debugger import set_trace","metadata":{"execution":{"iopub.status.busy":"2021-06-08T07:49:01.163313Z","iopub.execute_input":"2021-06-08T07:49:01.164110Z","iopub.status.idle":"2021-06-08T07:49:01.169637Z","shell.execute_reply.started":"2021-06-08T07:49:01.164057Z","shell.execute_reply":"2021-06-08T07:49:01.168409Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nfrom pprint import pprint\nfrom typing import Dict\n\nfrom tempfile import gettempdir\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport torch\nfrom torch import nn, optim\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader, random_split\nfrom torchvision.models.resnet import resnet50\n# from torchvision.models.mobilenet import mobilenet_v2\nfrom tqdm.notebook import tqdm\nimport wandb\n\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\nfrom prettytable import PrettyTable\nfrom pathlib import Path","metadata":{"execution":{"iopub.status.busy":"2021-06-08T07:58:07.787205Z","iopub.execute_input":"2021-06-08T07:58:07.787562Z","iopub.status.idle":"2021-06-08T07:58:07.796125Z","shell.execute_reply.started":"2021-06-08T07:58:07.787532Z","shell.execute_reply":"2021-06-08T07:58:07.795136Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Prepare Data path and load cfg\nBy setting the L5KIT_DATA_FOLDER variable, we can point the script to the folder where the data lies.\n\nThen, we load our config file with relative paths and other configurations (rasteriser, training params...).","metadata":{}},{"cell_type":"code","source":"# set env variable for data\nos.environ[\"L5KIT_DATA_FOLDER\"] = \"/kaggle/input/lyft-motion-prediction-autonomous-vehicles/\"\ndm = LocalDataManager(None)\n# get config\n# cfg = load_config_data(\"/kaggle/input/lyft-config-files/agent_motion_config.yaml\")\ncfg = load_config_data(\"/kaggle/input/kbl5kit-v2/examples/agent_motion_prediction/agent_motion_config.yaml\")\npprint(cfg)","metadata":{"execution":{"iopub.status.busy":"2021-06-08T07:49:05.788444Z","iopub.execute_input":"2021-06-08T07:49:05.788823Z","iopub.status.idle":"2021-06-08T07:49:05.817213Z","shell.execute_reply.started":"2021-06-08T07:49:05.788790Z","shell.execute_reply":"2021-06-08T07:49:05.816269Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cfg['model_params']['history_num_frames'] = 10\ncfg['train_params']['max_num_steps'] = 25000\ncfg['train_params']['checkpoint_every_n_steps'] = 5000\ncfg['train_data_loader']['batch_size'] = 12\ncfg['train_data_loader']['num_workers'] = 4\ncfg['train_data_loader']['key'] = 'scenes/train.zarr'\ncfg['val_data_loader']['batch_size'] = 12\ncfg['val_data_loader']['num_workers'] = 4\ncfg['val_data_loader']['key'] = 'scenes/validate.zarr'","metadata":{"execution":{"iopub.status.busy":"2021-06-08T07:49:08.590480Z","iopub.execute_input":"2021-06-08T07:49:08.590851Z","iopub.status.idle":"2021-06-08T07:49:08.597525Z","shell.execute_reply.started":"2021-06-08T07:49:08.590821Z","shell.execute_reply":"2021-06-08T07:49:08.596812Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model\n\nOur baseline is a simple `resnet50` pretrained on `imagenet`. We must replace the input and the final layer to address our requirements.","metadata":{}},{"cell_type":"code","source":"class Identity(nn.Module):\n    def __init__(self):\n        super(Identity, self).__init__()\n\n    def forward(self, x):\n        return x\n    \nclass Resnet50LSTM(nn.Module):\n    def __init__(self, config: Dict):\n        super(Resnet50LSTM, self).__init__()\n        self.cfg = config\n        self.batch_size = self.cfg['train_data_loader']['batch_size']\n        self.hist_frames = self.cfg['model_params']['history_num_frames']\n        self.num_targets = 2 * self.cfg[\"model_params\"][\"future_num_frames\"]\n        self.input_size = 8  # pos(2) + yaw(1) + vel(2) + accel(2) + yawrate(1)\n        self.hidden_size = 64\n        self.cnn = self.build_basecnn()\n        self.fc_infeatures = 2048 + self.hidden_size\n        self.fc0 = nn.Sequential(\n            nn.Linear(in_features=self.input_size, out_features=self.input_size),\n            nn.LeakyReLU(inplace=True)\n        )\n        self.fc1 = nn.Sequential(\n            nn.Dropout(p=0.2, inplace=False),\n            nn.Linear(in_features=self.fc_infeatures, out_features=4096),\n            nn.LeakyReLU(inplace=True),\n            nn.Dropout(p=0.2, inplace=False),\n            nn.Linear(in_features=4096, out_features=128),\n            nn.LeakyReLU(inplace=True)\n        )\n        self.encoder = nn.LSTM(input_size=self.input_size, hidden_size=self.hidden_size, num_layers=1, batch_first=True)\n        self.decoder = nn.LSTM(input_size=128, hidden_size=self.hidden_size, num_layers=1, batch_first=True)\n        self.fc2 = nn.Linear(in_features=(self.num_targets//2) * self.hidden_size, out_features=100)\n#         self.hidden_cell = (torch.zeros(self.batch_size, 1, self.hidden_size),\n#                             torch.zeros(self.batch_size, 1, self.hidden_size))\n\n    def forward(self, x, pos, yaw, vel, accel, yawrate):\n        enc_inputs = torch.cat([pos, yaw, vel, accel, yawrate], dim=2)\n        enc_inputs = [self.fc0(enc_inputs[:, i, :].reshape(-1, self.input_size)).unsqueeze(1) for i in range(enc_inputs.shape[1])]\n        enc_inputs = torch.cat(enc_inputs, dim=1)\n        enc_output, _ = self.encoder(enc_inputs)\n        enc_output = enc_output[:, -1, :]  # get ouput at last timestep\n        x = self.cnn(x)\n        x = torch.cat([x, enc_output], dim=1)\n        x = self.fc1(x)\n        x = torch.repeat_interleave(x.unsqueeze(1), repeats=self.num_targets//2, dim=1)  # shape: [batch, 50, 128]\n        dec_output, _ = self.decoder(x)  # shape: [batch, 50, 64]\n        x = self.fc2(dec_output.reshape(-1, (self.num_targets//2) * self.hidden_size))\n\n        return x\n\n    def build_basecnn(self):\n        # load pre-trained Conv2D model\n        model = resnet50(pretrained=True)\n\n        # change input channels number to match the rasterizer's output\n        num_history_channels = (cfg[\"model_params\"][\"history_num_frames\"] + 1) * 2\n        num_in_channels = 3 + num_history_channels\n        model.conv1 = nn.Conv2d(\n            num_in_channels,\n            model.conv1.out_channels,\n            kernel_size=model.conv1.kernel_size,\n            stride=model.conv1.stride,\n            padding=model.conv1.padding,\n            bias=False,\n        )\n\n        model.fc = Identity()\n    #     # change output size to (X, Y) * number of future states\n    #     num_targets = 2 * cfg[\"model_params\"][\"future_num_frames\"]\n    #     model.fc = nn.Linear(in_features=2048, out_features=num_targets)\n\n        return model","metadata":{"execution":{"iopub.status.busy":"2021-06-08T08:17:22.958947Z","iopub.execute_input":"2021-06-08T08:17:22.959323Z","iopub.status.idle":"2021-06-08T08:17:22.979012Z","shell.execute_reply.started":"2021-06-08T08:17:22.959293Z","shell.execute_reply":"2021-06-08T08:17:22.978139Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# class Identity(nn.Module):\n#     def __init__(self):\n#         super(Identity, self).__init__()\n\n#     def forward(self, x):\n#         return x\n    \n# class MobilenetV2LSTM(nn.Module):\n#     def __init__(self, config: Dict):\n#         super(MobilenetV2LSTM, self).__init__()\n#         self.cfg = config\n#         self.batch_size = self.cfg['train_data_loader']['batch_size']\n#         self.hist_frames = self.cfg['model_params']['history_num_frames']\n#         self.num_targets = 2 * self.cfg[\"model_params\"][\"future_num_frames\"]\n#         self.input_size = 8  # pos(2) + yaw(1) + vel(2) + accel(2) + yawrate(1)\n#         self.hidden_size = 64\n#         self.cnn = self.build_basecnn()\n#         self.fc_infeatures = 1280 + self.hidden_size\n#         self.fc0 = nn.Sequential(\n#             nn.Linear(in_features=self.input_size, out_features=self.input_size),\n#             nn.LeakyReLU(inplace=True)\n#         )\n#         self.fc1 = nn.Sequential(\n#             nn.Dropout(p=0.2, inplace=False),\n#             nn.Linear(in_features=self.fc_infeatures, out_features=4096),\n#             nn.LeakyReLU(inplace=True),\n#             nn.Dropout(p=0.2, inplace=False),\n#             nn.Linear(in_features=4096, out_features=128),\n#             nn.LeakyReLU(inplace=True)\n#         )\n#         self.encoder = nn.LSTM(input_size=self.input_size, hidden_size=self.hidden_size, num_layers=1, batch_first=True)\n#         self.decoder = nn.LSTM(input_size=128, hidden_size=self.hidden_size, num_layers=1, batch_first=True)\n#         self.fc2 = nn.Linear(in_features=(self.num_targets//2) * self.hidden_size, out_features=100)\n# #         self.hidden_cell = (torch.zeros(self.batch_size, 1, self.hidden_size),\n# #                             torch.zeros(self.batch_size, 1, self.hidden_size))\n\n#     def forward(self, x, pos, yaw, vel, accel, yawrate):\n#         enc_inputs = torch.cat([pos, yaw, vel, accel, yawrate], dim=2)\n#         enc_inputs = [self.fc0(enc_inputs[:, i, :].reshape(-1, self.input_size)).unsqueeze(1) for i in range(enc_inputs.shape[1])]\n#         enc_inputs = torch.cat(enc_inputs, dim=1)\n#         enc_output, _ = self.encoder(enc_inputs)\n#         enc_output = enc_output[:, -1, :]  # get ouput at last timestep\n#         x = self.cnn(x)\n#         x = torch.cat([x, enc_output], dim=1)\n#         x = self.fc1(x)\n#         x = torch.repeat_interleave(x.unsqueeze(1), repeats=self.num_targets//2, dim=1)  # shape: [batch, 50, 128]\n#         dec_output, _ = self.decoder(x)  # shape: [batch, 50, 64]\n#         x = self.fc2(dec_output.reshape(-1, (self.num_targets//2) * self.hidden_size))\n\n#         return x\n\n#     def build_basecnn(self):\n#         # change input channels number to match the rasterizer's output\n#         mnet = mobilenet_v2(pretrained=True)\n#         num_history_channels = (self.cfg[\"model_params\"][\"history_num_frames\"] + 1) * 2\n#         num_in_channels = 3 + num_history_channels\n#         mnet.features[0][0] = nn.Conv2d(\n#             num_in_channels,\n#             mnet.features[0][0].out_channels,\n#             kernel_size=mnet.features[0][0].kernel_size,\n#             stride=mnet.features[0][0].stride,\n#             padding=mnet.features[0][0].padding,\n#             bias=False,\n#         )\n\n#         mnet.classifier = Identity()\n        \n#         return mnet","metadata":{"execution":{"iopub.status.busy":"2021-06-06T19:14:07.725204Z","iopub.execute_input":"2021-06-06T19:14:07.725643Z","iopub.status.idle":"2021-06-06T19:14:07.968281Z","shell.execute_reply.started":"2021-06-06T19:14:07.725595Z","shell.execute_reply":"2021-06-06T19:14:07.966385Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def forward(data, model, device, criterion):\n    im_inputs = data[\"image\"].to(device)\n    # only (hist_frames - 2) accels available, so match all params to that\n    pos_inputs = data[\"history_positions\"][:, :-2, :].to(device)\n    yaw_inputs = data[\"history_yaws\"][:, :-2, :].to(device)\n    vel_inputs = data[\"history_velocities\"][:, :-1, :].to(device)\n    accel_inputs = data[\"history_accels\"].to(device)\n    yawrate_inputs = data[\"history_yawrates\"][:, :-1, :].to(device)\n    \n    # zero padding when there are not enough history frames in the data\n    if data['history_positions'].shape[1] < model.hist_frames + 1:\n        missing_frames = (model.hist_frames + 1) - data['history_positions'].shape[1]\n        pos_inputs = F.pad(pos_inputs, (0, 0, 0, missing_frames))\n        yaw_inputs = F.pad(yaw_inputs, (0, 0, 0, missing_frames))\n        vel_inputs = F.pad(vel_inputs, (0, 0, 0, missing_frames))\n        accel_inputs = F.pad(accel_inputs, (0, 0, 0, missing_frames))\n        yawrate_inputs = F.pad(yawrate_inputs, (0, 0, 0, missing_frames))\n\n    target_availabilities = data[\"target_availabilities\"].unsqueeze(-1).to(device)\n    targets = data[\"target_positions\"].to(device)\n    # Forward pass\n    outputs = model(im_inputs, pos_inputs, yaw_inputs, vel_inputs, accel_inputs, yawrate_inputs).reshape(targets.shape)\n    loss = criterion(outputs, targets)\n    # not all the output steps are valid, but we can filter them out from the loss using availabilities\n    loss = loss * target_availabilities\n    loss = loss.mean()\n    return loss, outputs","metadata":{"execution":{"iopub.status.busy":"2021-06-08T08:11:13.745948Z","iopub.execute_input":"2021-06-08T08:11:13.746333Z","iopub.status.idle":"2021-06-08T08:11:13.757813Z","shell.execute_reply.started":"2021-06-08T08:11:13.746299Z","shell.execute_reply":"2021-06-08T08:11:13.756637Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Load the Train Data\n\nOur data pipeline map a raw `.zarr` folder into a multi-processing instance ready for training by:\n- loading the `zarr` into a `ChunkedDataset` object. This object has a reference to the different arrays into the zarr (e.g. agents and traffic lights);\n- wrapping the `ChunkedDataset` into an `AgentDataset`, which inherits from torch `Dataset` class;\n- passing the `AgentDataset` into a torch `DataLoader`","metadata":{}},{"cell_type":"code","source":"# Load train dataset\ntrain_cfg = cfg[\"train_data_loader\"]\nrasterizer = build_rasterizer(cfg, dm)\ntrain_zarr = ChunkedDataset(dm.require(train_cfg[\"key\"])).open()\ntrain_dataset = AgentDataset(cfg, train_zarr, rasterizer)\ntrain_dataloader = DataLoader(train_dataset, shuffle=train_cfg[\"shuffle\"], batch_size=train_cfg[\"batch_size\"], \n                             num_workers=train_cfg[\"num_workers\"])\nprint(train_dataset)\nprint(len(train_dataset))\nprint(len(train_dataloader))","metadata":{"execution":{"iopub.status.busy":"2021-06-08T08:11:16.104360Z","iopub.execute_input":"2021-06-08T08:11:16.104741Z","iopub.status.idle":"2021-06-08T08:14:21.470228Z","shell.execute_reply.started":"2021-06-08T08:11:16.104693Z","shell.execute_reply":"2021-06-08T08:14:21.468629Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# ==== INIT MODEL\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n# model = MobilenetV2LSTM(cfg).to(device)\nmodel = Resnet50LSTM(cfg).to(device)\noptimizer = optim.AdamW(model.parameters(), lr=1e-5)\ncriterion = nn.MSELoss(reduction=\"none\")\nscheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=2500, eta_min=1e-8)","metadata":{"execution":{"iopub.status.busy":"2021-06-08T08:17:26.434233Z","iopub.execute_input":"2021-06-08T08:17:26.434751Z","iopub.status.idle":"2021-06-08T08:17:27.316643Z","shell.execute_reply.started":"2021-06-08T08:17:26.434705Z","shell.execute_reply":"2021-06-08T08:17:27.315604Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training loop","metadata":{}},{"cell_type":"code","source":"def train(model, device, train_loader, criterion, optimizer, scheduler, max_steps, checkpoint_steps):\n    # Switch model to training mode. This is necessary for layers like dropout, batchnorm etc \n    # which behave differently in training and evaluation mode\n    model.train()\n    losses_train = []\n    step = 0\n    \n    # We loop over the data iterator, and feed the inputs to the network and adjust the weights.\n    train_pbar = tqdm(enumerate(train_loader), desc=\"Training steps\", leave=True, total=max_steps)\n    for batch_idx, data in train_pbar:\n        # Forward pass to calculate loss\n        loss, _ = forward(data, model, device, criterion)\n\n        # Reset the gradients to 0 for all learnable weight parameters\n        optimizer.zero_grad()\n\n        # Backward pass: compute the gradients of the loss w.r.t. the model's parameters\n        loss.backward()\n\n        # Update the model weights\n        optimizer.step()\n        \n        scheduler.step()\n\n        # Get average loss of iterations so far\n        losses_train.append(loss.item())\n        avg_train_loss = np.mean(losses_train)\n\n        # wandb logging - metrics to track\n        wandb.log({'Training Loss': loss.item(), 'Avg Training Loss': avg_train_loss, \n                   'Learning rate': optimizer.param_groups[0][\"lr\"]})\n\n        train_pbar.set_description(f\" Avg train loss: {avg_train_loss}\")\n        \n#         if step % checkpoint_steps == 0:\n#             torch.save(model.state_dict(), f'l5run1_iter{step}_resnet50.pth')\n        \n        if step >= max_steps:\n            return avg_train_loss\n        step += 1\n\n    return avg_train_loss","metadata":{"execution":{"iopub.status.busy":"2021-06-08T08:17:30.531459Z","iopub.execute_input":"2021-06-08T08:17:30.531976Z","iopub.status.idle":"2021-06-08T08:17:30.539639Z","shell.execute_reply.started":"2021-06-08T08:17:30.531944Z","shell.execute_reply":"2021-06-08T08:17:30.538912Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Main","metadata":{}},{"cell_type":"code","source":"# Main training loop\n\n# Initialize wandb\nwandb.init(project='cs535-project', name='lstm-resnet50-25k', anonymous='must')\n\n# WandB – Config is a variable that holds and saves hyperparameters and inputs\nwandb_cfg = wandb.config                                         # initialize wandb config\nwandb_cfg.batch_size = cfg['train_data_loader']['batch_size']    # input batch size for training (default: 64)\nwandb_cfg.test_batch_size = cfg['val_data_loader']['batch_size'] # input batch size for testing (default: 1000)\nwandb_cfg.steps = cfg['train_params']['max_num_steps']\nwandb_cfg.epochs = 1                                             # number of epochs to train (default: 10)\nwandb_cfg.lr = 1e-5                                              # learning rate (default: 0.01)\nwandb_cfg.seed = 42                                              # random seed (default: 42)\nwandb_cfg.log_interval = 0                                       # how many batches to wait before logging training status\n\n# Set random seeds and deterministic pytorch for reproducibility\n# random.seed(wandb_cfg.seed)       # python random seed\n# torch.manual_seed(wandb_cfg.seed) # pytorch random seed\n# np.random.seed(wandb_cfg.seed)    # numpy random seed\n# torch.backends.cudnn.deterministic = True\n\n# WandB – wandb.watch() automatically fetches all layer dimensions, gradients, model parameters and logs them automatically to your dashboard.\n# Using log=\"all\" log histograms of parameter values in addition to gradients\n# wandb.watch(model)\n\n# pbar = tqdm(range(1, wandb_cfg.epochs+1), desc=\"Epochs\")\n# for epoch in pbar:\n#     train_loss = train(model, device, train_dataloader, criterion, optimizer, epoch)\n#     val_loss = val(model, device, val_dataloader, criterion, epoch)\n#     print(f\"Epoch: {epoch}, Train loss: {train_loss}, Val loss: {val_loss}\")\n\ncheckpoint_steps =  cfg['train_params']['checkpoint_every_n_steps']   \ntrain_loss = train(model, device, train_dataloader, criterion, optimizer, scheduler, wandb_cfg.steps, checkpoint_steps)","metadata":{"execution":{"iopub.status.busy":"2021-06-08T08:17:33.986580Z","iopub.execute_input":"2021-06-08T08:17:33.987132Z","iopub.status.idle":"2021-06-08T08:20:21.767381Z","shell.execute_reply.started":"2021-06-08T08:17:33.987100Z","shell.execute_reply":"2021-06-08T08:20:21.766388Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Save model\ntorch.save(model.state_dict(), 'cs535_lstm_resnet50.pth')","metadata":{},"execution_count":null,"outputs":[]}]}