{"cells":[{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import numpy as np\nimport os\nimport psutil\nimport torch\n\nfrom torch import nn\nfrom torch.utils.data import DataLoader\nfrom torchvision.models.mobilenet import mobilenet_v2\nfrom tqdm.notebook import tqdm\nfrom typing import Dict\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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"INPUT_DIR = '/kaggle/input/lyft-motion-prediction-autonomous-vehicles'\nWEIGHTS_FILE = '/kaggle/input/lyft-training-mobilenetv2/l5run3_mobilenetv2.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': [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    'test_data_loader': {\n        'key': 'scenes/test.zarr',\n        'batch_size': 8,\n        'shuffle': False,\n        'num_workers': 0\n    }\n\n}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# set env variable for data\nos.environ[\"L5KIT_DATA_FOLDER\"] = INPUT_DIR\ndm = LocalDataManager(None)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Init test dataset","execution_count":null},{"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\"{INPUT_DIR}/scenes/mask.npz\")[\"arr_0\"]\ntest_dataset = AgentDataset(cfg, test_zarr, rasterizer, agents_mask=test_mask)\n# test_dataset, _ = random_split(test_dataset, [100, 71122-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\"])\n\n\nprint(test_dataloader)\nprint(len(test_dataset))\nprint(len(test_dataloader))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Define model","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class Identity(nn.Module):\n    def __init__(self):\n        super(Identity, self).__init__()\n\n    def forward(self, x):\n        return x\n    \nclass MobilenetV2LSTM(nn.Module):\n    def __init__(self, config: Dict):\n        super(MobilenetV2LSTM, self).__init__()\n        self.cfg = config\n        self.batch_size = self.cfg['test_data_loader']['batch_size']\n        self.hist_frames = self.cfg['model_params']['history_num_frames']\n        self.fc_infeatures = 1280 + (2 * (self.hist_frames + 1)) + (2 * self.hist_frames) + (self.hist_frames + 1)\n        self.num_targets = 2 * self.cfg[\"model_params\"][\"future_num_frames\"]\n#         self.seq_len = 1\n#         self.input_size = 128\n#         self.hidden_size = 128\n        self.cnn = self.build_basecnn()\n        self.fc1 = nn.Sequential(\n            nn.Dropout(p=0.2, inplace=False),\n            nn.Linear(in_features=self.fc_infeatures, out_features=4096),\n            nn.ReLU(inplace=True),\n            nn.Dropout(p=0.2, inplace=False),\n            nn.Linear(in_features=4096, out_features=self.num_targets)\n        )\n#         self.lstm = nn.LSTM(\n#             input_size=self.input_size,\n#             hidden_size=self.hidden_size,\n#             num_layers=1,\n#             batch_first=True\n#         )\n#         self.fc2 = nn.Linear(in_features=128, out_features=100)\n#         self.hidden_cell = (torch.zeros(self.batch_size, 1, self.hidden_size),\n#                             torch.zeros(self.batch_size, 1, self.hidden_size))\n\n    def forward(self, x, vel, accel, yaw):\n        x = self.cnn(x)\n        vel = vel.reshape(self.batch_size, -1)\n        accel = accel.reshape(self.batch_size, -1)\n        yaw = yaw.reshape(self.batch_size, -1)\n        x = torch.cat([x, vel, accel, yaw], dim=1)\n        x = self.fc1(x)\n#         cnn_out = self.fc1(x)\n#         lstm_in = cnn_out.view(self.batch_size, self.seq_len, self.input_size)\n#         lstm_out, self.hidden_cell = self.lstm(lstm_in)\n#         fc_in = lstm_out.view(self.batch_size, lstm_out.shape[2])\n#         x = self.fc2(fc_in)\n\n        return x\n\n    def build_basecnn(self):\n        # change input channels number to match the rasterizer's output\n        mnet = mobilenet_v2(pretrained=False)\n        num_history_channels = (self.cfg[\"model_params\"][\"history_num_frames\"] + 1) * 2\n        num_in_channels = 3 + num_history_channels\n        mnet.features[0][0] = nn.Conv2d(\n            num_in_channels,\n            mnet.features[0][0].out_channels,\n            kernel_size=mnet.features[0][0].kernel_size,\n            stride=mnet.features[0][0].stride,\n            padding=mnet.features[0][0].padding,\n            bias=False,\n        )\n\n        mnet.classifier = Identity()\n        \n        return mnet","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# def build_model(cfg: Dict) -> torch.nn.Module:\n#     # load pre-trained Conv2D model\n#     model = mobilenet_v2(pretrained=False)\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.features[0][0] = nn.Conv2d(\n#         num_in_channels,\n#         model.features[0][0].out_channels,\n#         kernel_size=model.features[0][0].kernel_size,\n#         stride=model.features[0][0].stride,\n#         padding=model.features[0][0].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.classifier[1] = nn.Linear(in_features=model.classifier[1].in_features, out_features=num_targets)\n\n#     return model","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\")\nmodel = MobilenetV2LSTM(cfg).to(device)\n# model = build_model(cfg).to(device)\nmodel.load_state_dict(torch.load(WEIGHTS_FILE, map_location=device))\n# optimizer = optim.Adam(model.parameters(), lr=1e-3)\n# criterion = nn.MSELoss(reduction=\"none\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Prediction","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"model.eval()\n\nfuture_coords_offsets_pd = []\ntimestamps = []\nagent_ids = []\n\nwith torch.no_grad():\n    dataiter = iter(test_dataloader)\n    \n    pbar = tqdm(dataiter)\n    for data in pbar:\n        \n        im_inputs = data[\"image\"].to(device)\n        vel_inputs = data[\"history_velocities\"].to(device)\n        accel_inputs = data[\"history_accels\"][:, :-1, :].to(device)  # removing last history frame since we don't have accel for it\n        yaw_inputs = data[\"history_yaws\"].to(device)\n\n        target_availabilities = data[\"target_availabilities\"].unsqueeze(-1).to(device)\n        targets = data[\"target_positions\"].to(device)\n        # Forward pass\n        outputs = model(im_inputs, vel_inputs, accel_inputs, yaw_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        pbar.set_description(f'RAM used: {psutil.virtual_memory().percent}%')","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}