{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"import numpy as np\nimport torch\nimport os\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":{},"cell_type":"markdown","source":"# Configuration\n\n**The training (and EDA) part of this notebook can be found [here](https://www.kaggle.com/tuckerarrants/lyft-eda-training). In this notebook, we take our trained model and use it to predict on `test.zarr`. The model architecture is identical to the one  [Peter](https://www.kaggle.com/pestipeti) uses [here](https://www.kaggle.com/pestipeti/pytorch-baseline-train). The only difference is that the model was trained with `raster_size = [350, 350]`**","execution_count":null},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"cfg = {\n    'format_version': 4,\n    'model_params': {\n        'model_architecture': 'resnet18',\n        \n        'history_num_frames': 10,\n        'history_step_size': 1,\n        'history_delta_time': 0.1,\n        \n        'future_num_frames': 50,\n        'future_step_size': 1,\n        'future_delta_time': 0.1\n    },\n    \n    'raster_params': {\n        'raster_size': [350, 350],\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': 32,\n        'shuffle': True,\n        'num_workers': 0\n    },\n    \n    'val_data_loader': {\n        'key': 'scenes/validate.zarr',\n        'batch_size': 12,\n        'shuffle': False,\n        'num_workers': 4\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    'train_params': {\n        'checkpoint_every_n_steps': 5000,\n        'max_num_steps': 25000,\n        \n    }\n}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"PATH_TO_DATA = \"/kaggle/input/lyft-motion-prediction-autonomous-vehicles\"\nos.environ[\"L5KIT_DATA_FOLDER\"] = PATH_TO_DATA","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#get test.zarr into DataLoader form\ntest_cfg = cfg[\"test_data_loader\"]\ndm = LocalDataManager()\n\nrasterizer = build_rasterizer(cfg, dm)\n\ntest_zarr = ChunkedDataset(dm.require(test_cfg[\"key\"])).open()\ntest_mask = np.load(f\"{PATH_TO_DATA}/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_dataset)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model\n\n**The model below is taken straight from [Peter](https://www.kaggle.com/pestipeti)'s notebook [here](https://www.kaggle.com/pestipeti/pytorch-baseline-train):**","execution_count":null},{"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=False)\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        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        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":{},"cell_type":"markdown","source":"# Inference\n\n**Now we predict with the fully trained model defined above. Model has been trained for 25000 steps with a batch size of 32 and raster size [350, 350]**","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"#from training notebook\nPRE_TRAINED_MODEL_PATH = \"../input/lyftpretrainedmodels/model_state_24999.pth\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\nprint(device)\nmodel = LyftModel(cfg)\nmodel.to(device)\n\nmodel_state = torch.load(PRE_TRAINED_MODEL_PATH, map_location=device)\nmodel.load_state_dict(model_state)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"**This prediction loop is also taken from [Peter](https://www.kaggle.com/pestipeti)'s notebook [here](https://www.kaggle.com/pestipeti/pytorch-baseline-inference):**","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"#prediction loop\nmodel.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        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())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#save predictions as a csv file\nwrite_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}