{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"from typing import Dict\n\nfrom tempfile import gettempdir\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom torch import nn, optim\nfrom torch.utils.data import DataLoader\nimport torchvision\nfrom torchvision.models.resnet import resnet50, resnet18, resnet34, resnet101\nfrom tqdm import tqdm\n\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\nfrom prettytable import PrettyTable\nfrom pathlib import Path\n\nimport matplotlib.pyplot as plt\n\nimport os\nimport random\nimport time\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")\n\n\nfrom IPython.display import display\nfrom tqdm import tqdm_notebook\nimport gc, psutil\n\nprint(l5kit.__version__)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Memory measurement\ndef memory(verbose=True):\n    mem = psutil.virtual_memory()\n    gb = 1024*1024*1024\n    if verbose:\n        print('Physical memory:',\n              '%.2f GB (used),'%((mem.total - mem.available) / gb),\n              '%.2f GB (available)'%((mem.available) / gb), '/',\n              '%.2f GB'%(mem.total / gb))\n    return (mem.total - mem.available) / gb\n\ndef gc_memory(verbose=True):\n    m = gc.collect()\n    if verbose:\n        print('GC:', m, end=' | ')\n        memory()\n\nmemory();","execution_count":null,"outputs":[]},{"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)\nset_seed(42)","execution_count":null,"outputs":[]},{"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        'first_layer_bias': False,\n        'pretrained': True,\n        'multi_mode': True,\n        'model_architecture': 'resnet101',\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': \"model_resnet34_output\",\n        'lr': 1e-5,\n        'weight_path': '../input/resnet101-best-weights-epoch2/model_state_280000.pth',\n        'train': False,\n        'predict': True,\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    'train_data_loader': {\n        'key': 'scenes/train.zarr',\n        'batch_size': 16,\n        'shuffle': True,\n        'num_workers': 4,\n    },    \n    'test_data_loader': {\n        'key': 'scenes/test.zarr',\n        'batch_size': 32,\n        'shuffle': False,\n        'num_workers': 4,\n    },\n    'train_params': {\n#         'steps': 100,\n#         'update_steps': 10,\n#         'checkpoint_steps': 50,\n        'steps': 1000,\n        'update_steps': 100,\n        'checkpoint_steps': 1000,\n    },\n    \n    'combine': True\n    \n}\n\n# set env variable for data\nDIR_INPUT = cfg[\"data_path\"]\nos.environ[\"L5KIT_DATA_FOLDER\"] = DIR_INPUT\ndm = LocalDataManager()\n\n\n\n# Build rasterizer\nrasterizer = build_rasterizer(cfg, dm)\n\n\n# Test dataset\ntest_cfg = cfg[\"test_data_loader\"]\ntest_zarr = ChunkedDataset(dm.require(test_cfg[\"key\"])).open(cached=False)  # to prevent run out of memory\ntest_mask = np.load(f\"{DIR_INPUT}/scenes/mask.npz\")[\"arr_0\"]\ntest_dataset = AgentDataset(cfg, test_zarr, rasterizer, agents_mask=test_mask)\ntest_dataloader = DataLoader(test_dataset, shuffle=test_cfg[\"shuffle\"],\n                             batch_size=test_cfg[\"batch_size\"], num_workers=test_cfg[\"num_workers\"])\nprint(test_dataset)\n\n\n\n# --- 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)\n\n\nimport torch\nimport torchvision\n\nfrom torch import nn\nfrom torchvision.models.resnet import resnet50, resnet101, resnet34\nfrom typing import Dict\n\nclass LyftMultiModel(nn.Module):\n    def __init__(self, conf: Dict):\n        super().__init__()\n        \n        self.future_num_frames = cfg[\"model_params\"][\"future_num_frames\"]\n        target_count = 2 * self.future_num_frames\n\n        if cfg[\"model_params\"][\"multi_mode\"]:\n            self.multi_mode = True\n            target_count += 1   # One confidence per prediction\n            target_count *= 3   # 3 predictions instead of 1\n        else:\n            self.multi_mode = False\n\n        history_channel_count = (conf[\"model_params\"][\"history_num_frames\"] + 1) * 2\n        total_channel_count = 3 + history_channel_count\n\n        architecture = cfg[\"model_params\"][\"model_architecture\"]\n\n        if architecture == \"resnet50\":\n            backbone = resnet50(pretrained=cfg[\"model_params\"][\"pretrained\"])  \n        elif architecture == \"resnet101\":\n            backbone = resnet101(pretrained=cfg[\"model_params\"][\"pretrained\"])\n        elif architecture == \"resnet34\":\n            backbone = resnet34(pretrained=cfg[\"model_params\"][\"pretrained\"])\n            # architecture = conf[\"model_params\"][\"model_architecture\"]\n            # backbone = eval(architecture)(pretrained=True)\n\n        if architecture in [\"resnet50\", \"resnet101\", \"resnet34\"]:\n            backbone.conv1 = nn.Conv2d(\n                total_channel_count,\n                backbone.conv1.out_channels,\n                kernel_size=backbone.conv1.kernel_size,\n                stride=backbone.conv1.stride,\n                padding=backbone.conv1.padding,\n                bias=cfg[\"model_params\"][\"first_layer_bias\"], # Maybe True is better?\n            )\n\n            if architecture == \"resnet34\":\n                backbone.fc = nn.Linear(in_features=512, out_features=target_count)\n            else:\n                backbone.fc = nn.Linear(in_features=2048, out_features=target_count)\n\n            self.backbone = backbone\n\n    def forward(self, x):\n        y = self.backbone(x)\n\n        if self.multi_mode:\n            batches, _ = y.shape\n\n            # print(\"y.shape =\", y.shape)\n            # print(\"batches =\", batches)\n\n            pred, confidences = torch.split(y, self.future_num_frames * 3 * 2, dim=1)\n            pred = pred.view(batches, 3, self.future_num_frames, 2)\n\n            # print(\"confidences.shape = \", confidences.shape)\n\n            # assert confidences.shape == (batches, 3)\n\n            confidences = torch.softmax(confidences, dim=1)\n\n            return pred, confidences\n        else:\n            return y\n    \n    \ndef forward(data, model, device, criterion=pytorch_neg_multi_log_likelihood_batch, compute_loss=True):\n    inputs = data[\"image\"].to(device)\n    target_availabilities = data[\"target_availabilities\"].to(device)\n    targets = data[\"target_positions\"].to(device)\n    # Forward pass\n    preds, confidences = model(inputs)\n    # skip compute loss if we are doing prediction\n    loss = criterion(targets, preds, confidences, target_availabilities) if compute_loss else 0\n    return loss, preds, confidences\n\n\n\n\ndef execute(path):\n    # ==== INIT MODEL=================\n    device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n    model = LyftMultiModel(cfg)\n\n    #load weight if there is a pretrained model\n    weight_path = cfg[\"model_params\"][\"weight_path\"]\n    if weight_path:\n        model.load_state_dict(torch.load(weight_path))\n\n    model.to(device)\n    optimizer = optim.Adam(model.parameters(), lr=cfg[\"model_params\"][\"lr\"])\n    print(f'device {device}')\n\n\n\n\n    if cfg[\"model_params\"][\"predict\"]:\n\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        memorys_pred = []\n        t0 = time.time()\n        times_pred = []\n        iterations_pred = []\n\n        for i, data in enumerate(tqdm_notebook(test_dataloader, mininterval=5.)):\n\n            _, preds, confidences = forward(data, model, device, compute_loss=False)\n\n            preds = torch.einsum('bmti,bji->bmtj', \n                                 preds.double(), \n                                 data[\"world_from_agent\"].to(device)[:, :2, :2]).cpu().numpy()\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            if i%50 == 0:\n                t = ((time.time() - t0) / 60)\n                print('%4d'%i, '%6.2fmins'%t, end=' | ')\n                mem = memory()\n                iterations_pred.append(i)\n                memorys_pred.append(mem)\n                times_pred.append(t)\n    #             if i > 0:\n    #                 break\n        print('Total timespent: %6.2fmins'%((time.time() - t0) / 60))\n        memory()\n\n\n    # create submission to submit to Kaggle\n    pred_path = path\n    write_pred_csv(\n        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\n    df_sub = pd.read_csv(pred_path)\n    display(df_sub)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if cfg['combine']:\n\n    execute('submission_1.csv')\n\n    cfg[\"model_params\"][\"weight_path\"] = '../input/resnet34-bestscore-epoch-weights/model_state_last.pth'\n    cfg['model_params']['model_architecture'] = 'resnet34'\n\n    execute('submission_2.csv')\n    \nelse:\n    execute('submission.csv')\n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import pandas as pd\n\nif cfg['combine']:\n    pd.options.display.max_columns=305\n\n    paths = [\n        \"submission_1.csv\", \n        \"submission_2.csv\",\n    ]\n\n    weights = [0.2, 0.8]\n\n    conf_cols = np.array([\"conf_0\", \"conf_1\", \"conf_2\"])\n\n\n    xy_cols = [[],[],[]]\n    for i in range(50):\n        for j in range(3):\n            xy_cols[j].append(f\"coord_x{j}{i}\")\n            xy_cols[j].append(f\"coord_y{j}{i}\")\n    xy_cols[0][:10]\n\n\n    COLUMNS = [\"timestamp\", \"track_id\"] + list(conf_cols) + xy_cols[0] + xy_cols[1] + xy_cols[2]\n\n    def sort_df(df, sort_timestamp_track_id=True):\n\n        conf_orders = np.argsort(-df[conf_cols].values,1)\n        XY = np.stack([df[xy_cols[0]].values,df[xy_cols[1]].values, df[xy_cols[2]].values], axis=1)\n        XY = XY[np.arange(len(XY))[:, None], conf_orders]\n\n        df2 = pd.DataFrame(columns = COLUMNS)\n        df2[\"timestamp\"] = df[\"timestamp\"].values\n        df2[\"track_id\"] = df[\"track_id\"].values\n        df2[xy_cols[0] + xy_cols[1] + xy_cols[2]] = XY.reshape(-1,300)\n        df2[conf_cols] = df[conf_cols].values[np.arange(len(df))[:, None], conf_orders]\n\n        if sort_timestamp_track_id:\n            df2.sort_values([\"timestamp\", \"track_id\"], inplace=True)\n            df2.reset_index(inplace=True, drop=True)\n        return df2\n\n    df = None\n    for path,w in zip(paths,weights):\n        print(w, path)\n        temp = pd.read_csv(path)\n        temp = sort_df(temp)\n        temp[COLUMNS[5:]] *= w\n        if df is None:\n            df = temp\n        else:\n            df[COLUMNS[2:]] += temp[COLUMNS[2:]]\n    df[conf_cols] /= df[conf_cols].sum(1).values[:, None]\n\n    sample = pd.read_csv(\"../input/lyft-motion-prediction-autonomous-vehicles/multi_mode_sample_submission.csv\")\n\n    df = sample[[\"timestamp\", \"track_id\"]].merge(df, on=[\"timestamp\", \"track_id\"])\n    sample.shape, df.shape\n\n    df.to_csv(\"submission.csv\", index=False, float_format='%.6f')","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}