{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"collapsed":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":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"\nimport 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\nimport matplotlib.pyplot as plt\n\n\nfrom l5kit.data import LocalDataManager, ChunkedDataset\nfrom l5kit.dataset import AgentDataset, EgoDataset\nfrom l5kit.evaluation import write_pred_csv\nfrom l5kit.rasterization import build_rasterizer\n\nfrom l5kit.visualization import draw_trajectory, TARGET_POINTS_COLOR\nfrom l5kit.geometry import transform_points\n\nfrom torch.utils.data.dataset import Subset","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.cuda.is_available()","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\"\nDEBUG = True","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"cfg = {\n    'format_version': 4,\n    'model_params': {\n        'history_num_frames': 0,\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], # This is size of the image. That is 300*300 pixel\n        'pixel_size': [0.5, 0.5], # One pixel corresponds to these many meters. \n                                    #  So, along the width and height of the image, one pixel will correspond to 0.5 metres in the real world.\n        'ego_center': [0.25, 0.5], \n        'map_type': 'py_semantic', # The rasterizer to be used for visualization: \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': 8,\n        'shuffle': True,\n        'num_workers': 4\n    },\n    'test_data_loader': {\n        'key': 'scenes/test.zarr',\n        'batch_size': 16,\n        'shuffle': False,\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    'train_params': {\"epochs\": 10, 'checkpoint_every_n_steps': 200,\n          'max_num_steps':1000,\n          'eval_every_n_steps': 100\n    }\n\n}\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"os.environ[\"L5KIT_DATA_FOLDER\"] = DIR_INPUT\ndm = LocalDataManager(None)\n\ncfg[\"raster_params\"][\"map_type\"] = \"py_semantic\"\ncfg[\"model_params\"][\"history_num_frames\"] = 10","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Initiate Dataset\ntrain_cfg = cfg[\"train_data_loader\"]\n# Rasterizer\nrasterizer = build_rasterizer(cfg, dm)\n# Test dataset/dataloader\ntrain_zarr = ChunkedDataset(dm.require(train_cfg[\"key\"])).open()\n\ntrain_dataset = AgentDataset(cfg, train_zarr, rasterizer)\n\ntrain_dataset = Subset(train_dataset, np.arange(100))\n\ntrain_dataloader = DataLoader(train_dataset,\n                             shuffle=train_cfg[\"shuffle\"],\n                             batch_size=train_cfg[\"batch_size\"],\n                             num_workers=train_cfg[\"num_workers\"])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"collapsed":true},"cell_type":"code","source":"test_cfg = cfg[\"test_data_loader\"]\n\n# Rasterizer\n# rasterizer = 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)\n\n# test_dataset = Subset(test_dataset, np.arange(100))\ntest_dataloader = DataLoader(test_dataset,\n                             shuffle=test_cfg[\"shuffle\"],\n                             batch_size=test_cfg[\"batch_size\"],\n                             num_workers=test_cfg[\"num_workers\"])","execution_count":null,"outputs":[]},{"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        # X, Y coords for the future positions (output shape: Bx50x2)\n        num_targets = 2 * cfg[\"model_params\"][\"future_num_frames\"]\n\n        # You can add more layers here.\n        self.head = nn.Sequential(\n            # nn.Dropout(0.2),\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":"device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\nmodel = LyftModel(cfg)\nmodel.to(device)\noptimizer = optim.Adam(model.parameters(), lr=1e-3)\ncriterion = nn.MSELoss(reduction=\"none\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train"},{"metadata":{"trusted":true,"collapsed":true},"cell_type":"code","source":"# ==== TRAIN LOOP\ntr_it = iter(train_dataloader)\nprogress_bar = tqdm(range(cfg[\"train_params\"][\"max_num_steps\"]))\n\nlosses_train = []\nfor epoch in progress_bar:\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    lossees_train = []\n    inputs = data[\"image\"].to(device)\n    target_availabilities = data[\"target_availabilities\"].unsqueeze(-1).to(device)\n    targets = data[\"target_positions\"].to(device)\n    outputs = model(inputs).reshape(targets.shape)\n    loss = criterion(outputs, targets)\n    loss = loss * target_availabilities\n    loss = loss.mean()\n    \n    optimizer.zero_grad()\n    loss.backward()\n    optimizer.step()\n    losses_train.append(loss.item())\n    \n\n    if (epoch + 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)}\")\n    \n\n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Training Analysis\nplt.plot(losses_train, label=\"Training Loss\")\nplt.xlabel('Training step', fontsize=12) \nplt.ylabel('Loss', fontsize=12)\nplt.legend()\nplt.show()","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":{},"cell_type":"markdown","source":"# Predicting"},{"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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"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))","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}