{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Motion Prediction for Atonomous Vehicles\n\n**Notes/readMe**\n\n- `kaggle_l5kit` is used 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- In Kaggle notebook, all library/packages and datasets are accessible and imported from online sources, the code can be directly run in kaggle notebook.\n- To test the code, simply just click the \"Copy and Edit\" botton on the top right corner to open a new notebook for testing.\n- For local uses, the datasets need to be downloaded in the the same directory, Pytorch and l5kit packages need to be pre-installed into the environment. The code can be run directly in a IDE or use \"python fileName.py command in command window.\"\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\n\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","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\nDEBUG = True","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"cfg = {\n    'format_version': 4,\n    'model_params': {\n        'model_architecture': 'resnet50',\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    },\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    \n    'train_data_loader': {\n        'key': 'scenes/train.zarr',\n        'batch_size': 12,\n        'shuffle': True,\n        'num_workers': 4\n    },\n    \n    'train_params': {\n        'max_num_steps': 100 if DEBUG else 10000,\n        'checkpoint_every_n_steps': 5000,\n        \n        # 'eval_every_n_steps': -1\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":"## Dataset, dataloader for train set"},{"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# Train dataset/dataloader\ntrain_zarr = ChunkedDataset(dm.require(train_cfg[\"key\"])).open()\ntrain_dataset = AgentDataset(cfg, train_zarr, rasterizer)\ntrain_dataloader = DataLoader(train_dataset,\n                              shuffle=train_cfg[\"shuffle\"],\n                              batch_size=train_cfg[\"batch_size\"],\n                              num_workers=train_cfg[\"num_workers\"])\n\nprint(train_dataset)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"class LyftModel(nn.Module):\n    \n    def __init__(self, cfg: Dict):\n        super().__init__()\n        \n        self.backbone = resnet18(pretrained=True, progress=True)\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        # Predict X, Y coords for the future positions\n        num_targets = 2 * cfg[\"model_params\"][\"future_num_frames\"]\n\n        self.head = nn.Sequential(\n            \n            nn.Linear(in_features=backbone_out_features, out_features=4096),\n        )\n\n        self.logit = nn.Linear(4096, out_features=num_targets)\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        \n        return x","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\")\n\nmodel = LyftModel(cfg)\nmodel.to(device)\noptimizer = optim.Adam(model.parameters(), lr=1e-3)\n\n# Use to filter the invalid steps.\ncriterion = nn.MSELoss(reduction=\"none\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"# ==== TRAIN LOOP\ntr_it = iter(train_dataloader)\n\nprogress_bar = tqdm(range(cfg[\"train_params\"][\"max_num_steps\"]))\nlosses_train = []\n\nfor itr in progress_bar:\n\n    try:\n        data = next(tr_it)\n    except StopIteration:\n        tr_it = iter(train_dataloader)\n        data = next(tr_it)\n\n    model.train()\n    torch.set_grad_enabled(True)\n    \n    # Forward pass\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    outputs = model(inputs).reshape(targets.shape)\n    loss = criterion(outputs, targets)\n\n    # filter out valid output steps from the loss using availabilities\n    loss = loss * target_availabilities\n    loss = loss.mean()\n\n    # Backward pass\n    optimizer.zero_grad()\n    loss.backward()\n    optimizer.step()\n\n    losses_train.append(loss.item())\n\n    if (itr+1) % cfg['train_params']['checkpoint_every_n_steps'] == 0 and not DEBUG:\n        torch.save(model.state_dict(), f'model_state_{itr}.pth')\n    \n    progress_bar.set_description(f\"loss: {loss.item()} loss(avg): {np.mean(losses_train[-100:])}\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if not DEBUG:\n    torch.save(model.state_dict(), f'model_state_last.pth')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Load the output file from training\nWEIGHT_FILE = \"/kaggle/input/lyft-weights-private/exp_2229__iter_25000.pth\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"cfg = {\n    'format_version': 4,\n    'model_params': {\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    },\n    \n    'raster_params': {\n        'raster_size': [300, 300],\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    \n    'test_data_loader': {\n        'key': 'scenes/test.zarr',\n        'batch_size': 8,\n        'shuffle': False,\n        'num_workers': 4\n    }\n\n}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Dataset, dataloader for test set"},{"metadata":{"trusted":true},"cell_type":"code","source":"# ===== INIT DATASET\ntest_cfg = cfg[\"test_data_loader\"]\n\n# Rasterizer\nrasterizer = 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)\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":{"trusted":true},"cell_type":"code","source":"model.eval()\n\nfuture_coords_offsets_pd = []\ntimestamps = []\nagent_ids = []\n\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        outputs = model(inputs).reshape(targets.shape)\n        \n        future_coords_offsets_pd.append(outputs.cpu().numpy().copy())\n        timestamps.append(data[\"timestamp\"].numpy().copy())\n        agent_ids.append(data[\"track_id\"].numpy().copy())\n        \n        \n\nwrite_pred_csv('trajactoryPrediction.csv',\n               timestamps=np.concatenate(timestamps),\n               track_ids=np.concatenate(agent_ids),\n               coords=np.concatenate(future_coords_offsets_pd))","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}