{"cells":[{"metadata":{"id":"olZe0Y85doy-","trusted":true},"cell_type":"code","source":"\n!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n!python pytorch-xla-env-setup.py --version nightly --apt-packages libomp5 libopenblas-dev","execution_count":null,"outputs":[]},{"metadata":{"id":"G_nQ7piQdbN8","trusted":true},"cell_type":"code","source":"# this script transports l5kit and dependencies\n!pip -q install pymap3d==2.1.0 \n!pip -q install protobuf==3.12.2 \n!pip -q install transforms3d \n!pip -q install zarr \n!pip -q install ptable\n \n!pip -q install --no-dependencies l5kit","execution_count":null,"outputs":[]},{"metadata":{"id":"tMcOw0gqdfRi","trusted":true},"cell_type":"code","source":"\nimport numpy as np\nimport torch\nimport gc, os\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")\nimport numpy as np\nimport os\nimport torch\nfrom multiprocessing import Pool\nimport random\nimport bz2\nimport pickle\nfrom torch.nn import functional as f\n\n\nfrom torch import nn, optim\nfrom torch.utils.data import DataLoader,Dataset\nfrom torchvision.models.resnet import resnet18\nfrom tqdm import tqdm\nfrom typing import Dict,Tuple\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau,StepLR,ExponentialLR\n\n\nfrom l5kit.evaluation import write_pred_csv\nfrom l5kit.data import LocalDataManager\nfrom l5kit.data import LocalDataManager,filter_agents_by_labels,get_combined_scenes\nfrom l5kit.data import ChunkedDataset\nfrom l5kit.dataset import EgoDataset,AgentDataset\nfrom l5kit.rasterization import build_rasterizer\n\n\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,rmse,average_displacement_error_mean\n\n\n\nfrom pathlib import Path\nfrom tempfile import gettempdir\nimport pandas as pd\n\n\n\n","execution_count":null,"outputs":[]},{"metadata":{"id":"rbzamIebdrB1","trusted":true},"cell_type":"code","source":"!pip install pytorch-lightning","execution_count":null,"outputs":[]},{"metadata":{"id":"ZOoE4MgYdrpY","trusted":true},"cell_type":"code","source":"\nimport pytorch_lightning as pl\nfrom pytorch_lightning.loggers import CSVLogger","execution_count":null,"outputs":[]},{"metadata":{"id":"YXK_l0qUdxm-","trusted":true},"cell_type":"code","source":"\nDIR_INPUT = '../input/lyft-motion-prediction-autonomous-vehicles/' #data files\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 = False\nVALIDATION = False\n\n\n\n\n\ncfg = {\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': [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': 8,\n        'shuffle': True,\n        'num_workers': 4\n    },\n    \n     'valid_data_loader': {\n        'key': 'scenes/validate.zarr',\n        'batch_size': 8,\n        'shuffle': False,\n        'num_workers': 4\n    },\n    \n    'sample_data_loader': {\n        'key': 'scenes/sample.zarr',\n        'batch_size': 16,\n        'shuffle': False,\n        'num_workers': 0\n    },\n    \n    'train_params': {\n        'max_num_steps': 2334 if DEBUG else 20000,\n        'checkpoint_every_n_steps': 2000,\n        \n        # 'eval_every_n_steps': -1\n    }\n}\n\n\n# set env variable for data\nos.environ[\"L5KIT_DATA_FOLDER\"] = DIR_INPUT\ndm = LocalDataManager(None)\n\n\n\n","execution_count":null,"outputs":[]},{"metadata":{"id":"vpCNz3_1d48d","trusted":true},"cell_type":"code","source":"\n# --- Function utils ---\n# Original code from https://github.com/lyft/l5kit/blob/20ab033c01610d711c3d36e1963ecec86e8b85b6/l5kit/l5kit/evaluation/metrics.py\nimport numpy as np\n\nimport torch\nfrom torch import Tensor\n\n\ndef pytorch_neg_multi_log_likelihood_batch(\n    gt: Tensor, pred: Tensor, confidences: Tensor, avails: Tensor\n) -> Tensor:\n    \"\"\"\n    Compute a negative log-likelihood for the multi-modal scenario.\n    log-sum-exp trick is used here to avoid underflow and overflow, For more information about it see:\n    https://en.wikipedia.org/wiki/LogSumExp#log-sum-exp_trick_for_log-domain_calculations\n    https://timvieira.github.io/blog/post/2014/02/11/exp-normalize-trick/\n    https://leimao.github.io/blog/LogSumExp/\n    Args:\n        gt (Tensor): array of shape (bs)x(time)x(2D coords)\n        pred (Tensor): array of shape (bs)x(modes)x(time)x(2D coords)\n        confidences (Tensor): array of shape (bs)x(modes) with a confidence for each mode in each sample\n        avails (Tensor): array of shape (bs)x(time) with the availability for each gt timestep\n    Returns:\n        Tensor: negative log-likelihood for this example, a single float number\n    \"\"\"\n    assert len(pred.shape) == 4, f\"expected 3D (MxTxC) array for pred, got {pred.shape}\"\n    batch_size, num_modes, future_len, num_coords = pred.shape\n\n    assert gt.shape == (batch_size, future_len, num_coords), f\"expected 2D (Time x Coords) array for gt, got {gt.shape}\"\n    assert confidences.shape == (batch_size, num_modes), f\"expected 1D (Modes) array for gt, got {confidences.shape}\"\n    #print(confidences)\n    assert torch.allclose(torch.sum(confidences, dim=1), confidences.new_ones((batch_size,))), \"confidences should sum to 1\"\n    assert avails.shape == (batch_size, future_len), f\"expected 1D (Time) array for gt, got {avails.shape}\"\n    # assert all data are valid\n    assert torch.isfinite(pred).all(), \"invalid value found in pred\"\n    assert torch.isfinite(gt).all(), \"invalid value found in gt\"\n    assert torch.isfinite(confidences).all(), \"invalid value found in confidences\"\n    assert torch.isfinite(avails).all(), \"invalid value found in avails\"\n\n    # convert to (batch_size, num_modes, future_len, num_coords)\n    gt = torch.unsqueeze(gt, 1)  # add modes\n    avails = avails[:, None, :, None]  # add modes and cords\n\n    # error (batch_size, num_modes, future_len)\n    error = torch.sum(((gt - pred) * avails) ** 2, dim=-1)  # reduce coords and use availability\n\n    with np.errstate(divide=\"ignore\"):  # when confidence is 0 log goes to -inf, but we're fine with it\n        # error (batch_size, num_modes)\n        error = torch.log(confidences) - 0.5 * torch.sum(error, dim=-1)  # reduce time\n\n    # use max aggregator on modes for numerical stability\n    # error (batch_size, num_modes)\n    max_value, _ = error.max(dim=1, keepdim=True)  # error are negative at this point, so max() gives the minimum one\n    error = -torch.log(torch.sum(torch.exp(error - max_value), dim=-1, keepdim=True)) - max_value  # reduce modes\n    # print(\"error\", error)\n    return torch.mean(error)\n\n\ndef pytorch_neg_multi_log_likelihood_single(\n    gt: Tensor, pred: Tensor, avails: Tensor\n) -> Tensor:\n    \"\"\"\n\n    Args:\n        gt (Tensor): array of shape (bs)x(time)x(2D coords)\n        pred (Tensor): array of shape (bs)x(time)x(2D coords)\n        avails (Tensor): array of shape (bs)x(time) with the availability for each gt timestep\n    Returns:\n        Tensor: negative log-likelihood for this example, a single float number\n    \"\"\"\n    # pred (bs)x(time)x(2D coords) --> (bs)x(mode=1)x(time)x(2D coords)\n    # create confidence (bs)x(mode=1)\n    batch_size, future_len, num_coords = pred.shape\n    confidences = pred.new_ones((batch_size, 1))\n    return pytorch_neg_multi_log_likelihood_batch(gt, pred.unsqueeze(1), confidences, avails)","execution_count":null,"outputs":[]},{"metadata":{"id":"vf9D0hIFd5py","trusted":true},"cell_type":"code","source":"\nclass LyftModel(pl.LightningModule):\n    \"\"\"Model is resnet101_02 pretrained on imagenet.\n    We must replace the input and the final layer to address Lyft requirements.\n    \"\"\"\n\n    def __init__(self, cfg: Dict, pretrained=True):\n        super().__init__()\n\n        self.cfg = cfg\n        self.dm = LocalDataManager(None)\n        self.rast = build_rasterizer(self.cfg, self.dm)\n\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        self.future_len = cfg[\"model_params\"][\"future_num_frames\"]\n        num_targets = 2 * self.future_len\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+36, out_features=4096),\n        )\n        self.num_preds = num_targets * 3\n        self.num_modes = 3\n        \n        self.logit = nn.Linear(4096, out_features=self.num_preds + 3)\n\n\n    def forward(self, x,y):\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        #y = self.meta(y)\n        x = torch.cat((x, y), dim=1)\n        \n        x = self.head(x)\n        x = self.logit(x)\n        \n        bs, _ = x.shape\n        pred, confidences = torch.split(x, self.num_preds, dim=1)\n        pred = pred.view(bs, self.num_modes, self.future_len, 2)\n        assert confidences.shape == (bs, self.num_modes)\n        confidences = torch.softmax(confidences, dim=1)\n        \n        return pred, confidences\n\n\n    def training_step(self, batch, batch_idx):\n        target_availabilities = torch.tensor(\n            batch[\"target_availabilities\"], device=self.device\n        )\n        targets = torch.tensor(batch[\"target_positions\"], device=self.device)\n        data = torch.tensor(batch[\"image\"], device=self.device)\n        meta=torch.cat((batch['extent']\n                    ,torch.flatten(batch['history_yaws'].float(), 1)\n                    ,batch['history_positions'][:,:,0].float(),batch['history_positions'][:,:,1].float()), dim=1).to(self.device)\n\n\n        outputs,confidence = self(data,meta)\n        loss = pytorch_neg_multi_log_likelihood_batch(targets, outputs, confidence, target_availabilities)\n\n        pbar ={'train_loss':loss}\n        return {'loss':loss,'progress_bar':pbar}\n\n    def validation_step(self, batch, batch_idx):\n        target_availabilities = torch.tensor(\n            batch[\"target_availabilities\"], device=self.device\n        )\n        targets = torch.tensor(batch[\"target_positions\"], device=self.device)\n        data = torch.tensor(batch[\"image\"], device=self.device)\n        meta=torch.cat((batch['extent']\n                    ,torch.flatten(batch['history_yaws'].float(), 1)\n                    ,batch['history_positions'][:,:,0].float(),batch['history_positions'][:,:,1].float()), dim=1).to(self.device)\n\n\n        outputs,confidence = self(data,meta)\n        loss = pytorch_neg_multi_log_likelihood_batch(targets, outputs, confidence, target_availabilities)\n\n\n        return {'val_loss':loss}\n\n    def configure_optimizers(self):\n        return optim.Adam(self.parameters(), lr=1e-3)\n\n","execution_count":null,"outputs":[]},{"metadata":{"id":"bf1tJDV_d_xm","trusted":true},"cell_type":"code","source":"\n#===== 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_dataset = DataLoader(train_dataset,\n                              shuffle=train_cfg[\"shuffle\"],#shuffle=True\n                              batch_size=train_cfg[\"batch_size\"],#batch_size=24\n                              num_workers=train_cfg[\"num_workers\"])#num_workers=4\n\n\n\n#print(train_dataset)\n\nval_cfg = cfg[\"valid_data_loader\"]\n\n# Rasterizer\nrasterizer = build_rasterizer(cfg, dm)\n\n# Test dataset/dataloader\nval_zarr = ChunkedDataset(dm.require(val_cfg[\"key\"])).open()\n\nval_dataset = AgentDataset(cfg, val_zarr, rasterizer)\nval_dataset = DataLoader(val_dataset,\n                             shuffle=val_cfg[\"shuffle\"],\n                             batch_size=val_cfg[\"batch_size\"],\n                             num_workers=val_cfg[\"num_workers\"])\n\n\n#print(val_dataloader)\n\n","execution_count":null,"outputs":[]},{"metadata":{"id":"_68jHnkheI87","trusted":true},"cell_type":"code","source":"model = LyftModel(cfg,pretrained=True)\ntrainer = pl.Trainer(tpu_cores=1, max_steps=500)#,\n\ntrainer.fit(model,train_dataset,val_dataset) #problem\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","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}