{"cells":[{"metadata":{},"cell_type":"markdown","source":"## Install dependencies (`pip` and `l5kit`)"},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install -q --no-index -f ../input/kaggle-l5kit-110 pip\n!pip install -q --no-index -f ../input/kaggle-l5kit-110 l5kit","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Make sure we have the correct version"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import l5kit\nimport torch\nimport torchvision\nl5kit.__version__, torch.__version__, torchvision.__version__, torch.cuda.is_available()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"And also CUDA is still usable"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"assert torch.cuda.is_available(), \"must be on GPU to train\"\ntorch.device('cuda')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Configuration via YAML"},{"metadata":{"trusted":true},"cell_type":"code","source":"%%writefile agent_motion_config.yml\n# sample from\n# https://github.com/lyft/l5kit/blob/master/examples/agent_motion_prediction/agent_motion_config.yaml\n\n# Config format schema number\nformat_version: 4\n\n###################\n## Model options\nmodel_params:\n  model_architecture: \"resnet18\"\n\n  history_num_frames: 0\n  history_step_size: 1\n  history_delta_time: 0.1\n\n  future_num_frames: 50\n  future_step_size: 1\n  future_delta_time: 0.1\n\n###################\n## Input raster parameters\nraster_params:\n  # raster image size [pixels]\n  raster_size:\n    - 224\n    - 224\n  # raster's spatial resolution [meters per pixel]: the size in the real world one pixel corresponds to.\n  pixel_size:\n    - 0.5\n    - 0.5\n  # From 0 to 1 per axis, [0.5,0.5] would show the ego centered in the image.\n  ego_center:\n    - 0.25\n    - 0.5\n  map_type: \"py_semantic\"\n\n  # the keys are relative to the dataset environment variable\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\n  # e.g. 0.0 include every obstacle, 0.5 show those obstacles with >0.5 probability of being\n  # one of the classes we care about (cars, bikes, peds, etc.), >=1.0 filter all other agents.\n  filter_agents_threshold: 0.5\n\n  # whether to completely disable traffic light faces in the semantic rasterizer\n  disable_traffic_light_faces: False\n\n###################\n## Data loader options\ntrain_data_loader:\n  key: \"scenes/train.zarr\"\n  batch_size: 16\n  shuffle: True\n  num_workers: 4\n\nval_data_loader:\n  key: \"scenes/sample.zarr\"\n  batch_size: 16\n  shuffle: False\n  num_workers: 4\n\ntest_data_loader:\n  key: \"scenes/test.zarr\"\n  batch_size: 8\n  shuffle: False\n  num_workers: 4\n\n###################\n## Train params\ntrain_params:\n  checkpoint_every_n_steps: 10000\n  max_num_steps: 5\n  eval_every_n_steps: 10000\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Dataset definitions"},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nfrom l5kit.configs import load_config_data\nfrom l5kit.data import PERCEPTION_LABELS, ChunkedDataset, LocalDataManager\nfrom l5kit.dataset import AgentDataset, EgoDataset\nfrom l5kit.geometry import transform_points\nfrom l5kit.rasterization import build_rasterizer\nfrom l5kit.visualization import TARGET_POINTS_COLOR, draw_trajectory\nfrom prettytable import PrettyTable\nfrom torch.utils.data import DataLoader\n\nis_kaggle = os.path.isdir(\"/kaggle\")\n\nconfig_file = \"agent_motion_config.yml\"\n\ndata_root = (\n    \"/kaggle/input/lyft-motion-prediction-autonomous-vehicles\"\n    if is_kaggle\n    else \"lyft-motion-prediction-autonomous-vehicles\"\n)\n\n\nclass MyDataset(object):\n    def __init__(self, config_file: str = config_file, data_root: str = data_root):\n        super().__init__()\n        self.cfg = load_config_data(config_file)\n        self.dm = LocalDataManager(data_root)\n        self.rast = build_rasterizer(self.cfg, self.dm)\n        # self.ego_dataset = EgoDataset(self.cfg, self.zarr_dataset, self.rast)\n\n    def chunked_dataset(self, key: str):\n        dl_cfg = self.cfg[key]\n        dataset_path = self.dm.require(dl_cfg[\"key\"])\n        zarr_dataset = ChunkedDataset(dataset_path)\n        zarr_dataset.open()\n        return zarr_dataset\n\n    @property\n    def val_data_loader(self):\n        key = \"val_data_loader\"\n        dl_cfg = self.cfg[key]\n        zarr_dataset = self.chunked_dataset(key)\n        agent_dataset = AgentDataset(self.cfg, zarr_dataset, self.rast)\n        return DataLoader(\n            agent_dataset,\n            shuffle=dl_cfg[\"shuffle\"],\n            batch_size=dl_cfg[\"batch_size\"],\n            num_workers=dl_cfg[\"num_workers\"],\n            pin_memory=True,\n        )\n\n    @property\n    def test_data_loader(self):\n        key = \"test_data_loader\"\n        dl_cfg = self.cfg[key]\n        zarr_dataset = self.chunked_dataset(key)\n        test_mask = np.load(f\"{data_root}/scenes/mask.npz\")[\"arr_0\"]\n        agent_dataset = AgentDataset(\n            self.cfg, zarr_dataset, self.rast, agents_mask=test_mask\n        )\n        return DataLoader(\n            agent_dataset,\n            shuffle=dl_cfg[\"shuffle\"],\n            batch_size=dl_cfg[\"batch_size\"],\n            num_workers=dl_cfg[\"num_workers\"],\n            pin_memory=True,\n        )\n\n    @property\n    def train_data_loader(self):\n        key = \"train_data_loader\"\n        dl_cfg = self.cfg[key]\n        zarr_dataset = self.chunked_dataset(key)\n        agent_dataset = AgentDataset(self.cfg, zarr_dataset, self.rast)\n        return DataLoader(\n            agent_dataset,\n            shuffle=dl_cfg[\"shuffle\"],\n            batch_size=dl_cfg[\"batch_size\"],\n            num_workers=dl_cfg[\"num_workers\"],\n            pin_memory=True,\n        )\n\n    def plt_show_agent_map(self, idx):\n        data = self.agent_dataset[idx]\n        im = data[\"image\"].transpose(1, 2, 0)\n        im = self.rast.to_rgb(im)\n        target_positions_pixels = transform_points(\n            data[\"target_positions\"] + data[\"centroid\"][:2], data[\"world_to_image\"]\n        )\n        draw_trajectory(\n            im, target_positions_pixels, data[\"target_yaws\"], TARGET_POINTS_COLOR\n        )\n        plt.imshow(im[::-1])\n        plt.show()\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Model definition"},{"metadata":{"trusted":true},"cell_type":"code","source":"from typing import Dict\n\nimport numpy as np\nimport torch\nfrom l5kit.evaluation.csv_utils import write_pred_csv\nfrom l5kit.evaluation.metrics import neg_multi_log_likelihood\nfrom l5kit.geometry import transform_points\nfrom torch import nn, optim\nfrom torchvision.models.resnet import resnet18, resnet50\nfrom tqdm.notebook import tqdm\n\ntry:\n    from .data_loading import MyDataset\nexcept:\n    pass\n\nimport torch\nfrom torch import Tensor\n\n\n# loss definition from https://www.kaggle.com/corochann/lyft-training-with-multi-mode-confidence\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 == (\n        batch_size,\n        future_len,\n        num_coords,\n    ), f\"expected 2D (Time x Coords) array for gt, got {gt.shape}\"\n    assert confidences.shape == (\n        batch_size,\n        num_modes,\n    ), f\"expected 1D (Modes) array for confidences, got {confidences.shape}\"\n    assert torch.allclose(\n        torch.sum(confidences, dim=1), confidences.new_ones((batch_size,))\n    ), \"confidences should sum to 1\"\n    assert avails.shape == (\n        batch_size,\n        future_len,\n    ), f\"expected 1D (Time) array for avails, 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(\n        ((gt - pred) * avails) ** 2, dim=-1\n    )  # reduce coords and use availability\n\n    with np.errstate(\n        divide=\"ignore\"\n    ):  # 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(\n        dim=1, keepdim=True\n    )  # error are negative at this point, so max() gives the minimum one\n    error = (\n        -torch.log(torch.sum(torch.exp(error - max_value), dim=-1, keepdim=True))\n        - max_value\n    )  # reduce modes\n    # print(\"error\", error)\n    return torch.mean(error)\n\n\nclass BaselineModel(nn.Module):\n    \"\"\"Our baseline is a simple resnet34 pretrained on imagenet.\n    We must replace the input and the final layer to address our requirements.\n    \"\"\"\n\n    def __init__(self, cfg: Dict, num_modes: int = 3, pretrained=False):\n        super().__init__()\n        self.num_modes = num_modes\n        resnet = resnet50(pretrained=pretrained)\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        resnet.conv1 = nn.Conv2d(\n            num_in_channels,\n            resnet.conv1.out_channels,\n            kernel_size=resnet.conv1.kernel_size,\n            stride=resnet.conv1.stride,\n            padding=resnet.conv1.padding,\n            bias=resnet.conv1.bias,\n        )\n\n        # 512 for resnet18 and resnet34;\n        backbone_out_features = 2048\n\n        self.future_len = cfg[\"model_params\"][\"future_num_frames\"]\n        num_targets = 2 * self.future_len\n\n        self.num_preds = num_targets * self.num_modes\n\n        resnet.fc = nn.Sequential(\n            nn.Linear(\n                in_features=backbone_out_features,\n                # num of modes * preds + confidence\n                out_features=self.num_preds + self.num_modes,\n            ),\n        )\n        self.resnet = resnet\n\n    def forward(self, data):\n        out = self.resnet(data)\n        batch_size = data.shape[0]\n        pred, confidences = torch.split(out, self.num_preds, dim=1)\n        assert pred.shape == (batch_size, self.num_preds)\n        assert confidences.shape == (batch_size, self.num_modes)\n        pred = pred.view(batch_size, self.num_modes, self.future_len, 2)\n        confidences = torch.softmax(confidences, dim=1)\n        return pred, confidences\n\n\ndef training_step(model, criterion, device, batch):\n    target_availabilities = batch[\"target_availabilities\"].to(device)\n    targets = batch[\"target_positions\"].to(device)\n    data = batch[\"image\"].to(device)\n    pred, confidences = model(data)\n    return criterion(targets, pred, confidences, target_availabilities)\n\n\ndef train_one_epoch(epoch, model, criterion, device, optimizer, train_dl, max_steps):\n    for idx, batch in enumerate(tqdm(train_dl, desc=f\"Training Epoch {epoch}\")):\n        if idx > max_steps:\n            break\n        loss = training_step(\n            model=model,\n            criterion=criterion,\n            device=device,\n            batch=batch,\n        )\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n\n\ndef validation_step(model, criterion, device, batch):\n    target_availabilities = batch[\"target_availabilities\"].to(device)\n    targets = batch[\"target_positions\"].to(device)\n    data = batch[\"image\"].to(device)\n    pred, confidences = model(data)\n    return criterion(targets, pred, confidences, target_availabilities)\n\n\ndef validate_one_epoch(epoch, model, criterion, device, val_dl, max_steps: int):\n    all_loses = []\n    for idx, batch in enumerate(tqdm(val_dl, desc=f\"Validation Epoch {epoch}\")):\n        if idx > max_steps:\n            break\n        loss = validation_step(model, criterion, device=device, batch=batch)\n        all_loses.append(loss.mean())\n    return all_loses\n\n\ndef train_model(\n    dataset: MyDataset,\n    max_epochs=5,\n    max_steps=3000,\n    max_val_steps=100,\n    learning_rate=1e-3,\n    pretrained=False,\n):\n    assert torch.cuda.is_available(), \"gpu must be used\"\n    device = torch.device(\"cuda\")\n    print(\"training model using\", device)\n    model = BaselineModel(dataset.cfg, pretrained=pretrained).to(device)\n    train_dl = dataset.train_data_loader\n    val_dl = dataset.val_data_loader\n    criterion = pytorch_neg_multi_log_likelihood_batch\n    optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n    for epoch in tqdm(range(max_epochs), desc=\"Training\"):\n        train_one_epoch(\n            epoch,\n            model=model,\n            criterion=criterion,\n            optimizer=optimizer,\n            device=device,\n            train_dl=train_dl,\n            max_steps=max_steps,\n        )\n\n        with torch.no_grad():\n            validate_one_epoch(\n                epoch=epoch,\n                model=model,\n                criterion=criterion,\n                device=device,\n                val_dl=val_dl,\n                max_steps=max_val_steps,\n            )\n    return model\n\n\ndef convert_agent_coordinates_to_world_offsets(\n    agents_coords: np.ndarray,\n    world_from_agents: np.ndarray,\n    centroids: np.ndarray,\n) -> np.ndarray:\n    coords_offset = []\n    for agent_coords, world_from_agent, centroid in zip(\n        agents_coords, world_from_agents, centroids\n    ):\n        predition_offset = []\n        for agent_coord in agent_coords:\n            predition_offset.append(\n                transform_points(agent_coord, world_from_agent) - centroid[:2]\n            )\n        predition_offset = np.stack(predition_offset)\n        coords_offset.append(predition_offset)\n    return np.stack(coords_offset)\n\n\ndef evaluation(cfg: Dict, model_path: str, dataset: MyDataset):\n    assert torch.cuda.is_available(), \"GPU must be used\"\n    device = torch.device(\"cuda\")\n\n    model = BaselineModel(cfg=cfg)\n    model.load_state_dict(torch.load(model_path))\n    model = model.to(device)\n    model.eval()\n\n    with torch.no_grad():\n        pred_coords_list = []\n        confidences_list = []\n        timestamps_list = []\n        track_id_list = []\n\n        for data in tqdm(dataset.test_data_loader):\n            inputs = data[\"image\"].to(device)\n            # target_availabilities = data[\"target_availabilities\"].to(device)\n            # targets = data[\"target_positions\"].to(device)\n            pred, confidences = model(inputs)\n            pred = convert_agent_coordinates_to_world_offsets(\n                pred.detach().cpu().numpy(),\n                data[\"world_from_agent\"].numpy(),\n                data[\"centroid\"].numpy(),\n            )\n            pred_coords_list.append(pred)\n            confidences_list.append(confidences.detach().cpu().numpy())\n            timestamps_list.append(data[\"timestamp\"].detach().numpy())\n            track_id_list.append(data[\"track_id\"].detach().numpy())\n\n        timestamps = np.concatenate(timestamps_list)\n        track_ids = np.concatenate(track_id_list)\n        coords = np.concatenate(pred_coords_list)\n        confs = np.concatenate(confidences_list)\n\n    return timestamps, track_ids, coords, confs\n\n\nif __name__ == \"__main__\":\n    print(\"loading dataset\")\n    my_dataset = MyDataset()\n    print(\"dataset is loaded, starting training\")\n    model = train_model(my_dataset)\n    print(\"training done, saving state and generating results\")\n    torch.save(model.state_dict(), \"model_state_last.pth\")\n\n    timestamps, track_ids, coords, confs = evaluation(\n        my_dataset.cfg, \"model_state_last.pth\", my_dataset\n    )\n    write_pred_csv(\n        \"submission.csv\",\n        timestamps=timestamps,\n        track_ids=track_ids,\n        coords=coords,\n        confs=confs,\n    )\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}