{"cells":[{"metadata":{},"cell_type":"markdown","source":"### Baseline Model\n\nThis baseline model is adopted from [Lyft's example](https://github.com/lyft/l5kit/blob/master/examples/agent_motion_prediction/agent_motion_prediction.ipynb) on their l5kit repo."},{"metadata":{},"cell_type":"markdown","source":"### Installing l5kit"},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"!pip install --upgrade pip\n!pip install pymap3d==2.1.0\n!pip install -U l5kit","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Importing PyTorch and l5kit"},{"metadata":{"trusted":true},"cell_type":"code","source":" !pip3 install resnet_pytorch","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","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 torch\nfrom torch import nn, optim\nfrom torch.utils.data import DataLoader\n\nfrom tqdm import tqdm\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\n\nimport os","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from torchvision.models.resnet import resnet18","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Prepare data path and config file"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true,"_kg_hide-output":false},"cell_type":"code","source":"# set env variable for data\nos.environ[\"L5KIT_DATA_FOLDER\"] = \"../input/lyft-motion-prediction-autonomous-vehicles\"\ndm = LocalDataManager(None)\n# get config\n#cfg = load_config_data(\"../input/lyft-config-files/agent_motion_config.yaml\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#MODEL_NAME = \"wide_resnet18\"\nIMG_SIZE = 224\n# --- Lyft configs ---\ncfg = {\n          'model_params': {'model_architecture': 'resnet18',\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_resnet101_output\",\n        'lr': 1e-4,\n        'weight_path': \"/kaggle/input/lyftpretrained-resnet101/lyft_resnet101_model.pth\",\n        'train': True,\n        'predict': True},\n\n        'raster_params': {'raster_size': [IMG_SIZE, IMG_SIZE],\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': {'key': 'scenes/train.zarr',\n          'batch_size': 8,\n          'shuffle': True,\n          'num_workers': 0},\n\n        \"valid_data_loader\":{\"key\": \"scenes/validate.zarr\",\n                            \"batch_size\": 8,\n                            \"shuffle\": False,\n                            \"num_workers\": 0},\n    \n        \"sample_data_loader\": {\n        'key': 'scenes/sample.zarr',\n        'batch_size': 8,\n        'shuffle': False,\n        'num_workers': 0},\n         \n        \"test_data_loader\":{\n        'key': \"scenes/test.zarr\",\n        'batch_size': 8,\n        'shuffle': False,\n        'num_workers': 0},\n\n    \n        'train_params': {\"epochs\": 10, 'checkpoint_every_n_steps': 200,\n          'max_num_steps':1000,\n          'eval_every_n_steps': 100}\n        }\nprint(cfg)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Config is where you can make your changes to have different `model_architecture`, `history_step_size`, `history_num_frames`, `batch_size`, etc. Inspect `cfg` for more details."},{"metadata":{},"cell_type":"markdown","source":"### Build a baseline CNN with Resnet50 backbone\n\nSize of `num_targets`: 100"},{"metadata":{"trusted":true},"cell_type":"code","source":"def build_model(cfg: Dict) -> torch.nn.Module:\n    # load pre-trained Conv2D model\n    model = resnet18(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    # 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=512, out_features=num_targets)\n\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def forward(data, model, device, criterion):\n    inputs = data[\"image\"].to(device)\n    target_availabilities = data[\"target_availabilities\"].unsqueeze(-1).to(device)\n    targets = data[\"target_positions\"].to(device)\n    # Forward pass\n    outputs = model(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","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Loading the data"},{"metadata":{},"cell_type":"markdown","source":"In the config file the `train_data_loader`'s key is a sample zarr file. Change this to `train.zarr` file either by doing something like below or chaning in the config file itself. If you are using kaggle's GPU, you can increase the batch size too. The default batch is 16 and it only takes around 2GB of GPU memory while you train. The number of workers to load the data is set to 16. You can reduce this a bit to put less work on the CPU."},{"metadata":{"trusted":true},"cell_type":"code","source":"# ===== INIT 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)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# ===== INIT VALIDATION DATASET============================================================\nvalid_cfg = cfg[\"valid_data_loader\"]\nrasterizer = build_rasterizer(cfg, dm)\nvalidate_zarr = ChunkedDataset(dm.require(valid_cfg[\"key\"])).open()\nvalid_dataset = AgentDataset(cfg, validate_zarr, rasterizer)\nvalid_dataloader = DataLoader(valid_dataset, shuffle=valid_cfg[\"shuffle\"], batch_size=valid_cfg[\"batch_size\"], \n                             num_workers=valid_cfg[\"num_workers\"])\nprint(\"==================================VALIDATION DATA==================================\")\nprint(valid_dataset)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Build model, set optimizer and loss function"},{"metadata":{"trusted":true},"cell_type":"code","source":"# ==== INIT MODEL\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\nmodel = build_model(cfg).to(device)\noptimizer = optim.Adam(model.parameters(), lr=1e-4)\ncriterion = nn.MSELoss(reduction=\"none\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Train model\n\nTrain for 1000 steps. "},{"metadata":{"trusted":true},"cell_type":"code","source":"VALIDATION = True","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train(train_dataloader, valid_dataloader, opt=None, criterion=None, lrate=1e-4):\n        \"\"\"Function for training the model\"\"\"\n        print(\"Building Model...\")\n        device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n        model = build_model(cfg).to(device)\n        optimizer = optim.Adam(model.parameters(), lr=1e-4)\n        criterion = nn.MSELoss(reduction=\"none\")\n                             \n                \n        print(\"Training...\")\n        losses = []\n        losses_mean = []\n        \n        val_losses = []\n        val_losses_mean = []\n        \n        progress = tqdm(range(cfg[\"train_params\"][\"max_num_steps\"]))\n        \n        train_iter = iter(train_dataloader)\n        val_iter = iter(valid_dataloader)\n        \n        for i in progress:\n            try:\n                data = next(train_iter)\n            except StopIteration:\n                train_iter = iter(train_dataloader)\n                data = next(train_iter)\n                    \n            model.train()\n            torch.set_grad_enabled(True)\n                    \n            loss, _ = forward(data, model, device, criterion)\n                        \n                    \n            optimizer.zero_grad()\n            loss.backward()\n            optimizer.step()\n            \n            # Validation\n            if VALIDATION:\n                with torch.no_grad():\n                    try:\n                        val_data = next(val_iter)\n                    except StopIteration:\n                        val_iter = iter(val_dataloader)\n                        val_data = next(val_iter)\n\n                    val_loss, _  = forward(val_data, model, device, criterion)\n                    val_losses.append(val_loss.item())\n                    val_losses_mean.append(np.mean(val_losses))\n                    \n                desc = f\"Loss: {round(loss.item(), 4)} Validation Loss: {round(val_loss.item(), 4)}\"\n            else:\n                desc = f\"Loss: {round(loss.item(), 4)}\"\n                \n            #if len(losses)>0 and loss < min(losses):\n            #    print(f\"Loss improved from {min(losses)} to {loss}\")\n                \n            \n            \n            losses.append(loss.item())\n            losses_mean.append(np.mean(losses))\n            progress.set_description(desc)\n            \n        return losses_mean, val_losses_mean, model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"losses, val_losses, model = train(train_dataloader, valid_dataloader)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Training Analysis\nplt.plot(losses, c=\"red\", label=\"Mean Training Loss\")\nplt.plot(val_losses, c=\"green\", label=\"Mean Validation Loss\")\nplt.xlabel('Training step', fontsize=12) \nplt.ylabel('Loss', fontsize=12)\nplt.legend()\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Load evaluation dataset"},{"metadata":{},"cell_type":"markdown","source":"Due to the fact that the following steps take way too long, they are commented out."},{"metadata":{"trusted":true},"cell_type":"code","source":"# # Loading eval dataset\neval_cfg = cfg[\"sample_data_loader\"]\nrasterizer = build_rasterizer(cfg, dm)\neval_zarr = ChunkedDataset(dm.require(eval_cfg[\"key\"])).open()\neval_dataset = AgentDataset(cfg, eval_zarr, rasterizer)\neval_dataloader = DataLoader(eval_dataset, \n                             shuffle=eval_cfg[\"shuffle\"], \n                             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":"#==== EVAL LOOP\nmodel.eval()\ntorch.set_grad_enabled(False)\n# store information for evaluation\nfuture_coords_offsets_pd = []\ntimestamps = []\n\nagent_ids = []\nprogress_bar = tqdm(eval_dataloader)\nfor data in progress_bar:\n    _, ouputs = forward(data, model, device, criterion)\n    future_coords_offsets_pd.append(ouputs.cpu().numpy().copy())\n    timestamps.append(data[\"timestamp\"].numpy().copy())\n    agent_ids.append(data[\"track_id\"].numpy().copy())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"error = compute_error_csv(eval_gt_path, pred_path)\nprint(f\"NLL: {error:.5f}\\nL2: {np.sqrt(2*error/cfg['model_params']['future_num_frames']):.5f}\")","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}