{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"### Setup L5Kit\nDataset Link: *https://www.kaggle.com/rhtsingh/kaggle-l5kit*","metadata":{}},{"cell_type":"code","source":"!pip install --no-index -f ../input/kaggle-l5kit pip==20.2.2 >/dev/nul\n!pip install --no-index -f ../input/kaggle-l5kit -U l5kit > /dev/nul","metadata":{"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Setup Dependencies\n\n*References: https://github.com/lyft/l5kit/blob/master/examples/agent_motion_prediction/agent_motion_prediction.ipynb*","metadata":{}},{"cell_type":"code","source":"import os\nimport gc\ngc.enable()\n\nimport numpy as np\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision\nfrom torch.utils.data import DataLoader\n\nfrom tqdm import tqdm \nfrom typing import Dict\nfrom pathlib import Path\nfrom prettytable import PrettyTable\n\nfrom l5kit.configs import load_config_data\nfrom l5kit.data import LocalDataManager, ChunkedDataset\nfrom l5kit.dataset import EgoDataset, AgentDataset\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","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Prepare Data path and load cfg\n\nBy setting the L5KIT_DATA_FOLDER variable, we can point the script to the folder where the data lies.\n\nThen, we load our config file with relative paths and other configurations (rasteriser, training params...).","metadata":{}},{"cell_type":"code","source":"DIR_INPUT = '../input/lyft-motion-prediction-autonomous-vehicles/'\nDIR_INPUT_TRAIN = '../input/lyft-full-training-set/'\nSINGLE_MODE_SUBMISSION = f\"{DIR_INPUT}/single_mode_submission.csv\"\nMULTI_MODE_SUBMISSION = f\"{DIR_INPUT}/multi_mode_submission.csv\"\n\nDEBUG = False\n\nos.environ['L5KIT_DATA_FOLDER'] = DIR_INPUT\ndm = LocalDataManager(None)\n\ncfg = {\n    'format_version':4,\n    'model_params':{\n        'model_architecture':'resnet18',\n        'history_num_frames':15,\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    'raster_params':{\n        'raster_size':[331,331],\n        'pixel_size':[0.5,0.5],\n        'ego_center':[0.25,0.25],\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    'train_params':{\n        'max_num_steps': 1000 if DEBUG else 20000,\n        'checkpoint_every_n_steps':5000\n    },\n    'test_data_loader': {\n        'key': 'scenes/test.zarr',\n        'batch_size': 8,\n        'shuffle': False,\n        'num_workers': 4\n    }\n}","metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Load the Data\n\nOur data pipeline map a raw .zarr folder into a multi-processing instance ready for training by:\n\n* loading the `zarr` into a ChunkedDataset object. This object has a reference to the different arrays into the zarr (e.g. agents and traffic lights);\n* wrapping the ChunkedDataset into an AgentDataset, which inherits from torch Dataset class;\n* pa*ssing the AgentDataset into a torch DataLoader","metadata":{}},{"cell_type":"code","source":"train_cfg = cfg['train_data_loader']\nrasterizer = build_rasterizer(cfg, dm)\n\ntrain_zarr = ChunkedDataset(dm.require(train_cfg['key'])).open()\ntrain_dataset = AgentDataset(cfg, train_zarr, rasterizer)\ntrain_dataloader = DataLoader(\n    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)","metadata":{"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Model\n*References: https://www.kaggle.com/pestipeti/pytorch-baseline-train*\n\nOur baseline is a simple resnet50 pretrained on imagenet. We must replace the input and the final layer to address our requirements.","metadata":{}},{"cell_type":"code","source":"class ResNetModel(nn.Module):\n    def __init__(self, cfg):\n        super(ResNetModel, self).__init__()\n        # load pre-trained Conv2D model\n        self.backbone = torchvision.models.resnet18(pretrained=False, progress=False)\n        self.backbone.load_state_dict(\n            torch.load(\n                '../input/resnet18/resnet18.pth'\n            )\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        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        # change output size to (X, Y) * number of future states\n        num_targets = 2 * cfg[\"model_params\"][\"future_num_frames\"]\n        self.backbone.fc = nn.Linear(in_features=512, out_features=num_targets)        \n    \n    def forward(self, x):\n        # Forward pass\n        return self.backbone(x)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Training Config\n\nHere we set up the training configuration i.e. device, laoding model, optimizer and loss function.","metadata":{}},{"cell_type":"code","source":"device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\nmodel = ResNetModel(cfg).to(device)\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\ncriterion = nn.MSELoss(reduction=\"none\")\nprint(\"Device is {}.\".format(device))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Training\n\n*note: if you're on MacOS and using py_satellite rasterizer, you may need to disable opencv multiprocessing by adding: cv2.setNumThreads(0) before the following cell. This seems to only affect running in python notebook and it's caused by the cv2.warpaffine function*","metadata":{}},{"cell_type":"code","source":"train_iter = iter(train_dataloader)\n\nprogress_bar = tqdm(range(cfg[\"train_params\"][\"max_num_steps\"]))\nlosses_train = []\nfor itr in progress_bar:\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    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    loss = loss * target_availabilities # filter invalid steps\n    loss = loss.mean()\n    optimizer.zero_grad()\n    loss.backward()\n    optimizer.step()\n\n    losses_train.append(loss.item())\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:])}\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"del train_dataloader, train_dataset, train_iter, train_zarr, model\ngc.collect()","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Evaluation","metadata":{"trusted":true}},{"cell_type":"code","source":"# WEIGHT_FILE = 'model_state_9999.pth'\n\n# test_cfg = cfg[\"test_data_loader\"]\n\n# # Rasterizer\n# rasterizer = build_rasterizer(cfg, dm)\n\n# # Test dataset/dataloader\n# test_zarr = ChunkedDataset(dm.require(test_cfg[\"key\"])).open()\n# test_mask = np.load(f\"{DIR_INPUT}/scenes/mask.npz\")[\"arr_0\"]\n# test_dataset = AgentDataset(cfg, test_zarr, rasterizer, agents_mask=test_mask)\n# test_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\n# print(test_dataloader)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# model = ResNetModel(cfg=cfg).to(device)\n# if WEIGHT_FILE is not None:\n#     model.load_state_dict(\n#         torch.load(WEIGHT_FILE, map_location=device),\n#     )\n# print(f\"Running on {device}.\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# model.eval()\n\n# future_coords_offsets_pd = []\n# timestamps = []\n# agent_ids = []\n\n# with 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())","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Save Results\nAfter the model has predicted trajectories for our evaluation set, we can save them in a csv file.","metadata":{}},{"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))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}