{"cells":[{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","collapsed":true,"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":false},"cell_type":"markdown","source":"# Pytorch Baseline\n\n**Notes**\n- Do not forget to enable the GPU (TPU) for training\n- You have to add `kaggle_l5kit` as utility script\n- Parts of the code below is from the [official example](https://github.com/lyft/l5kit/blob/master/examples/agent_motion_prediction/agent_motion_prediction.ipynb)\n"},{"metadata":{"trusted":true},"cell_type":"code","source":"import numpy as np\nimport os\nimport torch\n\nfrom torch import nn, optim\nfrom torch.utils.data import DataLoader\nfrom torchvision.models.resnet import resnet18\nfrom tqdm import tqdm\nfrom typing import Dict\nimport matplotlib.pyplot as plt\n\n\nfrom l5kit.data import LocalDataManager, ChunkedDataset\nfrom l5kit.dataset import AgentDataset, EgoDataset\nfrom l5kit.evaluation import write_pred_csv\nfrom l5kit.rasterization import build_rasterizer\n\nfrom l5kit.visualization import draw_trajectory, TARGET_POINTS_COLOR\nfrom l5kit.geometry import transform_points","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from torch.utils.data.dataset import Subset\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"DIR_INPUT = \"/kaggle/input/lyft-motion-prediction-autonomous-vehicles\"\n\nSINGLE_MODE_SUBMISSION = f\"{DIR_INPUT}/single_mode_sample_submission.csv\"\nMULTI_MODE_SUBMISSION = f\"{DIR_INPUT}/multi_mode_sample_submission.csv\"\n\n# Training notebook's output.\nWEIGHT_FILE = None","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Config"},{"metadata":{"trusted":true},"cell_type":"code","source":"cfg = {\n    'format_version': 4,\n    'model_params': {\n        'history_num_frames': 0,\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    },\n    \n    'raster_params': {\n        'raster_size': [300, 300], # This is size of the image. That is 300*300 pixel\n        'pixel_size': [0.5, 0.5], # One pixel corresponds to these many meters. \n                                    #  So, along the width and height of the image, one pixel will correspond to 0.5 metres in the real world.\n        'ego_center': [0.25, 0.5], \n        'map_type': 'py_semantic', # The rasterizer to be used for visualization: \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    'train_data_loader': {\n        'key': 'scenes/train.zarr',\n        'batch_size': 8,\n        'shuffle': True,\n        'num_workers': 4\n    },\n    'test_data_loader': {\n        'key': 'scenes/test.zarr',\n        'batch_size': 16,\n        'shuffle': False,\n        'num_workers': 4\n    }\n\n}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# set env variable for data\nos.environ[\"L5KIT_DATA_FOLDER\"] = DIR_INPUT\ndm = LocalDataManager(None)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Visualize"},{"metadata":{"trusted":true},"cell_type":"code","source":"zarr_dataset = ChunkedDataset('../input/lyft-motion-prediction-autonomous-vehicles/scenes/sample.zarr')\nzarr_dataset.open()\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"rast = build_rasterizer(cfg, dm)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Visulaize AV"},{"metadata":{},"cell_type":"markdown","source":"`EgoDataset`, `AgentDataset`   and  will allow us to iterate over the autnomous vehicle annotations."},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset_ego = EgoDataset(cfg, zarr_dataset, rast)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"data = dataset_ego[0]\n\nim = data[\"image\"].transpose(1, 2, 0)\nim = dataset_ego.rasterizer.to_rgb(im)\ntarget_positions_pixels = transform_points(data[\"target_positions\"] + data[\"centroid\"][:2], data[\"world_to_image\"])\ndraw_trajectory(im, target_positions_pixels, data[\"target_yaws\"], TARGET_POINTS_COLOR)\n\nplt.imshow(im[::-1])\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"This image contain:\n* The green box: It is the autnomous vehicle.\n* The pink line: The trajectory of the automous vehicle.\n* The blue box: An agent.\n* The yellow lines: The semantic maps.\n* But what are the orange outlined boxes: Most prbably traffic posts, but not veru sure."},{"metadata":{},"cell_type":"markdown","source":"## Visualize Agent"},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset_agent = AgentDataset(cfg, zarr_dataset, rast)\ndata = dataset_agent[0]\n\nim = data[\"image\"].transpose(1, 2, 0)\nim = dataset_agent.rasterizer.to_rgb(im)\ntarget_positions_pixels = transform_points(data[\"target_positions\"] + data[\"centroid\"][:2], data[\"world_to_image\"])\ndraw_trajectory(im, target_positions_pixels, data[\"target_yaws\"], TARGET_POINTS_COLOR)\n\nplt.imshow(im[::-1])\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Visualize satellite image"},{"metadata":{"trusted":true},"cell_type":"code","source":"# map_type was changed from 'py_semantic' to 'py_satellite'.\ncfg[\"raster_params\"][\"map_type\"] = \"py_satellite\"\nsatellite_rasterizer = build_rasterizer(cfg, dm)\nsatellite_dataset = EgoDataset(cfg, zarr_dataset, satellite_rasterizer)\n\ndata = satellite_dataset[0]\nim = data[\"image\"].transpose(1, 2, 0)\nim = dataset_agent.rasterizer.to_rgb(im)\nplt.imshow(im[::-1])\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Dataset, dataloader"},{"metadata":{},"cell_type":"markdown","source":"Change config to load train data.\n\nUse 10 frames history for model"},{"metadata":{"trusted":true},"cell_type":"code","source":"cfg[\"raster_params\"][\"map_type\"] = \"py_semantic\"\ncfg[\"model_params\"][\"history_num_frames\"] = 10\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# ===== INIT DATASET\ntrain_cfg = cfg[\"train_data_loader\"]\n\n# Rasterizer\nrasterizer = build_rasterizer(cfg, dm)\n\n# Test dataset/dataloader\ntrain_zarr = ChunkedDataset(dm.require(train_cfg[\"key\"])).open()\n# train_mask = np.load(f\"{DIR_INPUT}/scenes/mask.npz\")[\"arr_0\"]\ntrain_dataset = AgentDataset(cfg, train_zarr, rasterizer)\n\ntrain_dataset = Subset(train_dataset, np.arange(100))\n\ntrain_dataloader = DataLoader(train_dataset,\n                             shuffle=train_cfg[\"shuffle\"],\n                             batch_size=train_cfg[\"batch_size\"],\n                             num_workers=train_cfg[\"num_workers\"])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# ===== INIT DATASET\ntest_cfg = cfg[\"test_data_loader\"]\n\n# Rasterizer\n# rasterizer = build_rasterizer(cfg, dm)\n\n# Test dataset/dataloader\ntest_zarr = ChunkedDataset(dm.require(test_cfg[\"key\"])).open()\ntest_mask = np.load(f\"{DIR_INPUT}/scenes/mask.npz\")[\"arr_0\"]\ntest_dataset = AgentDataset(cfg, test_zarr, rasterizer, agents_mask=test_mask)\n\n# test_dataset = Subset(test_dataset, np.arange(100))\ntest_dataloader = DataLoader(test_dataset,\n                             shuffle=test_cfg[\"shuffle\"],\n                             batch_size=test_cfg[\"batch_size\"],\n                             num_workers=test_cfg[\"num_workers\"])\n\n\nprint(test_dataloader)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Shape of input image"},{"metadata":{"trusted":true},"cell_type":"code","source":"test_dataset[0]['image'].shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_dataset[101]['track_id']","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_dataset[0]['target_positions'].shape","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Model"},{"metadata":{},"cell_type":"markdown","source":"Model:\n* input : channel * image_size * image_size . Ex with 10 history num frame and image size = 300*300. Input = (25*300*300)\n* output: num_node + num_node * 2 * future_num_frames . Ex with 3-mode and 50 frame predict. Output have shape = (303,)"},{"metadata":{"trusted":true},"cell_type":"code","source":"class LyftModel(nn.Module):\n    \n    def __init__(self, cfg: Dict,num_modes = 3):\n        super().__init__()\n        \n        self.backbone = resnet18(pretrained=False)\n        \n        num_history_channels = (cfg[\"model_params\"][\"history_num_frames\"] + 1) * 2\n        num_in_channels = 3 + num_history_channels\n\n        self.backbone.conv1 = nn.Conv2d(\n            num_in_channels,\n            self.backbone.conv1.out_channels,\n            kernel_size=self.backbone.conv1.kernel_size,\n            stride=self.backbone.conv1.stride,\n            padding=self.backbone.conv1.padding,\n            bias=False,\n        )\n        \n        # This is 512 for resnet18 and resnet34;\n        # And it is 2048 for the other resnets\n        backbone_out_features = 512\n\n        # X, Y coords for the future positions (output shape: Bx50x2)\n        self.future_len = cfg[\"model_params\"][\"future_num_frames\"]\n        num_targets = 2 * self.future_len\n\n        # You can add more layers here.\n        self.head = nn.Sequential(\n            # nn.Dropout(0.2),\n            nn.Linear(in_features=backbone_out_features, out_features=4096),\n        )\n        self.num_preds = num_targets * num_modes\n        self.num_modes = num_modes\n        self.logit = nn.Linear(4096, out_features=self.num_preds+self.num_modes)\n        \n    def forward(self, x):\n        x = self.backbone.conv1(x)\n        x = self.backbone.bn1(x)\n        x = self.backbone.relu(x)\n        x = self.backbone.maxpool(x)\n\n        x = self.backbone.layer1(x)\n        x = self.backbone.layer2(x)\n        x = self.backbone.layer3(x)\n        x = self.backbone.layer4(x)\n\n        x = self.backbone.avgpool(x)\n        x = torch.flatten(x, 1)\n        \n        x = self.head(x)\n        x = self.logit(x)\n        bs, _ = x.shape\n        pred, confidences = torch.split(x, self.num_preds, dim=1)\n        pred = pred.view(bs, self.num_modes, self.future_len, 2)\n        assert confidences.shape == (bs, self.num_modes)\n        confidences = torch.softmax(confidences, dim=1)\n        return pred, confidences\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# ==== INIT MODEL\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\nmodel = LyftModel(cfg,num_modes=3)\nmodel.to(device)\n\nif WEIGHT_FILE is not None:\n    # Saved state dict from the training notebook\n    model_state = torch.load(WEIGHT_FILE, map_location=device)\n    model.load_state_dict(model_state['model_state_dict'])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training"},{"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\nimport numpy as np\n\nimport torch\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":"## Train"},{"metadata":{"trusted":true},"cell_type":"code","source":"optimizer = optim.Adam(model.parameters(), lr=1e-3)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model.train()\nfor epoch in range(1):\n    running_loss = 0.0\n    for data in tqdm(train_dataloader):\n        inputs = data[\"image\"].to(device)\n        target_availabilities = data[\"target_availabilities\"].to(device)\n        targets = data[\"target_positions\"].to(device)\n        optimizer.zero_grad()\n        pred, confidences  = model(inputs)\n        loss = pytorch_neg_multi_log_likelihood_batch(targets, pred, confidences, target_availabilities)\n        loss.backward()\n        optimizer.step()\n        running_loss += loss.item() * inputs.size(0)\n    print(running_loss)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Predicting"},{"metadata":{"trusted":true},"cell_type":"code","source":"model.eval()\n\nfuture_coords_offsets_pd = []\ntimestamps = []\nagent_ids = []\nconfidences_list = []\nwith torch.no_grad():\n    dataiter = tqdm(test_dataloader)\n    \n    for data in dataiter:\n\n        inputs = data[\"image\"].to(device)\n        target_availabilities = data[\"target_availabilities\"].unsqueeze(-1).to(device)\n        targets = data[\"target_positions\"].to(device)\n\n        pred, confidences  = model(inputs)\n        \n        future_coords_offsets_pd.append(pred.cpu().numpy().copy())\n        timestamps.append(data[\"timestamp\"].numpy().copy())\n        agent_ids.append(data[\"track_id\"].numpy().copy())\n        confidences_list.append(confidences.cpu().numpy().copy())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"write_pred_csv('submission.csv',\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              )","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}