{"cells":[{"metadata":{},"cell_type":"markdown","source":"### Baseline Model\n\nThis baseline model is adopted from [Lyft's example](https://github.com/lyft/l5kit/blob/master/examples/agent_motion_prediction/agent_motion_prediction.ipynb) on their l5kit repo."},{"metadata":{"trusted":true},"cell_type":"code","source":"!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":{"trusted":true},"cell_type":"code","source":"import warnings\nimport torch_xla\nimport torch_xla.debug.metrics as met\nimport torch_xla.distributed.data_parallel as dp\nimport torch_xla.distributed.parallel_loader as tlp\nimport torch_xla.utils.utils as xu\nimport torch_xla.core.xla_model as xm\nimport torch_xla.distributed.xla_multiprocessing as xmp\nimport torch_xla.test.test_utils as test_utils\nimport warnings","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Installing l5kit"},{"metadata":{"trusted":true,"_kg_hide-output":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":{},"cell_type":"markdown","source":"### Importing PyTorch and l5kit"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"from typing import Dict\n\nfrom tempfile import gettempdir\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport torch\nfrom torch import nn, optim\nfrom torch.utils.data import DataLoader\nfrom tqdm import tqdm\n\nfrom l5kit.configs import load_config_data\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, 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\nfrom prettytable import PrettyTable\nfrom pathlib import Path\n\nimport os","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Prepare data path and config file"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true,"_kg_hide-output":false},"cell_type":"code","source":"# set env variable for data\nos.environ[\"L5KIT_DATA_FOLDER\"] = \"../input/lyft-motion-prediction-autonomous-vehicles\"\ndm = LocalDataManager(None)\n# get config\ncfg = {\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': [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    'train_data_loader': {\n        'key': 'scenes/train.zarr',\n        'batch_size': 4,\n        'shuffle': True,\n        'num_workers': 4\n    },\n    \n    'test_data_loader': {\n        'key': 'scenes/test.zarr',\n        'batch_size': 12,\n        'shuffle': False,\n        'num_workers': 4\n    }\n\n}\n\nTRAIN_MODE = False","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Config is where you can make your changes to have different `model_architecture`, `history_step_size`, `history_num_frames`, `batch_size`, etc. Inspect `cfg` for more details."},{"metadata":{"trusted":true},"cell_type":"code","source":"# ===== INIT DATASET\nif TRAIN_MODE:\n    train_cfg = cfg[\"train_data_loader\"]\n    rasterizer = build_rasterizer(cfg, dm)\n    train_zarr = ChunkedDataset(dm.require(train_cfg[\"key\"])).open()\n    train_dataset = AgentDataset(cfg, train_zarr, rasterizer)\n    train_dataloader = DataLoader(train_dataset, shuffle=train_cfg[\"shuffle\"], \n                                  batch_size=train_cfg[\"batch_size\"], \n                                 num_workers=train_cfg[\"num_workers\"])\n    print(train_dataset)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Model Define: Efficientnet-b0 + L2norm + Binary Head"},{"metadata":{"trusted":true},"cell_type":"code","source":"import sys\nsys.path.insert(0, \"../input/effnetpt\")\nfrom efficientnet_pytorch import EfficientNet\nimport torch\nimport torch.nn as nn\n\n\ndef l2_norm(input, axis=1):\n    norm = torch.norm(input, 2, axis, True)\n    output = torch.div(input, norm)\n    return output\n\n\nclass BinaryHead(nn.Module):\n    def __init__(self, num_class=4, emb_size=2048, s=16.0):\n        super(BinaryHead, self).__init__()\n        self.s = s\n        self.fc = nn.Sequential(nn.Linear(emb_size, num_class))\n\n    def forward(self, fea):\n        fea = l2_norm(fea)\n        logit = self.fc(fea) * self.s\n        return logit\n\n\nclass LyftModel(nn.Module):\n    def __init__(self, cfg: Dict):\n        super(LyftModel, self).__init__()\n\n        self.backbone = EfficientNet.from_name(\"efficientnet-b1\")\n        num_history_channels = (cfg[\"model_params\"][\"history_num_frames\"] + 1) * 2\n        num_in_channels = 3 + num_history_channels\n        num_targets = 2 * cfg[\"model_params\"][\"future_num_frames\"]\n        self.backbone._conv_stem = nn.Conv2d(\n            num_in_channels,\n            self.backbone._conv_stem.out_channels,\n            kernel_size=self.backbone._conv_stem.kernel_size,\n            stride=self.backbone._conv_stem.stride,\n            padding=self.backbone._conv_stem.padding,\n            bias=False,\n        )\n        self.avg_pool = nn.AdaptiveAvgPool2d((1, 1))\n        self.fea_bn = nn.BatchNorm1d(1280)\n        self.fea_bn.bias.requires_grad_(False)\n        self.binary_head = BinaryHead(num_targets, emb_size=1280, s=1)\n        self.dropout = nn.Dropout(p=0.2)\n \n\n    def forward(self, x):\n\n        img_feature = self.backbone.extract_features(x)\n        img_feature = self.avg_pool(img_feature)\n        img_feature = img_feature.view(img_feature.size(0), -1)\n        fea = self.fea_bn(img_feature)\n        # fea = self.dropout(fea)\n        output = self.binary_head(fea)\n\n        return output","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 = LyftModel(cfg)\n#     ckpt = torch.load('../input/lyftmodelall/effnet0l2binay_368.bin')\n#     model.load_state_dict(ckpt)\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"build_model(cfg)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Loading the data"},{"metadata":{"trusted":true},"cell_type":"code","source":"class AverageMeter(object):\n    def __init__(self, name, fmt=':f'):\n        self.name = name\n        self.fmt = fmt\n        self.reset()\n\n    def reset(self):\n        self.val = 0\n        self.avg = 0\n        self.sum = 0\n        self.count = 0\n\n    def update(self, val, n=1):\n        self.val = val\n        self.sum += val * n\n        self.count += n\n        self.avg = self.sum / self.count\n\n    def __str__(self):\n        fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'\n        return fmtstr.format(**self.__dict__)\n\nclass ProgressMeter(object):\n    def __init__(self, num_batches, meters, prefix=\"\"):\n        self.batch_fmtstr = self._get_batch_fmtstr(num_batches)\n        self.meters = meters\n        self.prefix = prefix\n\n    def display(self, batch):\n        entries = [self.prefix + self.batch_fmtstr.format(batch)]\n        entries += [str(meter) for meter in self.meters]\n        print('\\t'.join(entries))\n\n    def _get_batch_fmtstr(self, num_batches):\n        num_digits = len(str(num_batches // 1))\n        fmt = '{:' + str(num_digits) + 'd}'\n        return '[' + fmt + '/' + fmt.format(num_batches) + ']'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_loop_fn(train_loader, model, optimizer, device, scheduler, epoch=None):\n    # Train\n    batch_time = AverageMeter('Time', ':6.1f')\n    data_time = AverageMeter('Data', ':6.1f')\n    losses = AverageMeter('Loss', ':6.1e')\n    progress = ProgressMeter(\n        len(train_loader),\n        [batch_time, data_time, losses],\n        prefix=\"[xla:{}]Train:  Epoch: [{}]\".format(xm.get_ordinal(), epoch)\n    )\n    criterion = nn.MSELoss(reduction=\"none\")\n    model.train()\n    end = time.time()\n    for i, data in enumerate(train_loader):\n        data_time.update(time.time()-end)\n        inputs = data[\"image\"].to(device)\n        target_availabilities = data[\"target_availabilities\"].unsqueeze(-1).to(device)\n        targets = data[\"target_positions\"].to(device)\n        optimizer.zero_grad()\n        # Forward pass\n        outputs = model(inputs).reshape(targets.shape)\n        loss = criterion(outputs, targets)\n        # not all the output steps are valid, but we can filter them out from the loss using availabilities\n        loss = loss * target_availabilities\n        loss = loss.mean()\n        loss.backward()\n        xm.optimizer_step(optimizer)\n        losses.update(loss.item(), inputs.size(0))\n        scheduler.step(metrics=loss)\n        batch_time.update(time.time() - end)\n        end = time.time()\n        if i % 50 == 0:\n            progress.display(i)\n            xm.save(model.state_dict(), \"model.bin\")\n    del loss\n    del outputs\n    gc.collect()\n    \n    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from torch.utils.data.distributed import DistributedSampler\nimport torch_xla.debug.metrics as met\nWRAPPED_MODEL = xmp.MpModelWrapper(build_model(cfg))\n\ndef _run():\n    TRAIN_BATCH_SIZE = 8\n    EPOCHS = 1\n    xm.master_print('Starting Run ...')\n    train_sampler = DistributedSampler(\n        train_dataset,\n        num_replicas=xm.xrt_world_size(),\n        rank=xm.get_ordinal(),\n        shuffle=False\n    )\n    \n    train_data_loader = DataLoader(\n        train_dataset,\n        batch_size=TRAIN_BATCH_SIZE,\n        sampler=train_sampler,\n        drop_last=True,\n        num_workers=0\n    )\n    xm.master_print('Train Loader Created.')\n    \n    num_train_steps = int(len(train_dataset) / TRAIN_BATCH_SIZE / xm.xrt_world_size())\n    device = xm.xla_device()\n    model = WRAPPED_MODEL.to(device)\n    xm.master_print('Done Model Loading.')\n    optimizer = optim.AdamW(model.parameters(), lr=1e-4)\n    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min',\n                                                           factor=0.8, patience=3, \n                                                           verbose=False, \n                                                           threshold=0.0001, min_lr=0.00001)\n    xm.master_print(f'Num Train Steps= {num_train_steps}, XRT World Size= {xm.xrt_world_size()}.')\n    \n    for epoch in range(EPOCHS):\n        para_loader = tlp.ParallelLoader(train_data_loader, [device])\n        xm.master_print('Parallel Loader Created. Training ...')\n        train_loop_fn(para_loader.per_device_loader(device),\n                      model,  \n                      optimizer, \n                      device, \n                      scheduler, \n                      epoch\n                     )\n        \n        xm.master_print(\"Finished training epoch {}\".format(epoch))\n        if epoch == EPOCHS-1:\n            xm.master_print('Saving Model ..')\n            xm.save(model.state_dict(), \"model.bin\")\n            xm.master_print('Model Saved.')\n            \n    if METRICS_DEBUG:\n      xm.master_print(met.metrics_report(), flush=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import time\nfrom torch.nn import functional as F\ndef _mp_fn(rank, flags):\n#     torch.set_default_tensor_type('torch.FloatTensor')\n    _run()\n\nFLAGS={}\nif TRAIN_MODE:\n    xmp.spawn(_mp_fn, args=(FLAGS,), nprocs=8, start_method='fork')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### TEST MODE"},{"metadata":{},"cell_type":"markdown","source":"Due to the fact that the following steps take way too long, they are commented out."},{"metadata":{"trusted":true},"cell_type":"code","source":"import numpy as np\nimport os\nimport torch\n\nfrom torch import nn, optim\nfrom torch.utils.data import DataLoader\nfrom tqdm 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":"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\"\n\n# Training notebook's output.\nWEIGHT_FILE = \"../input/lyftmodelall/resnet50binaryhead.bin\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = LyftModel(cfg)\nckpt = torch.load('../input/lyftmodelall/effnet0l2binay_368.bin')\nmodel.load_state_dict(ckpt)","execution_count":null,"outputs":[]},{"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\"{DIR_INPUT}/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_dataloader)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def _test():\n    \n\n    future_coords_offsets_pd = []\n    timestamps = []\n    agent_ids = []\n    device = 'xla:0'\n    print(f\"device: {device} ready!\")\n    model = LyftModel(cfg)\n    ckpt = torch.load('../input/lyftmodelall/effnet0l2binay_368.bin')\n    model.load_state_dict(ckpt)\n    model = model.to(device)\n    model.eval()\n    with torch.no_grad():\n        dataiter = tqdm(test_dataloader)\n\n        for data in dataiter:\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    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":{"trusted":true},"cell_type":"code","source":"def _mp_fn(rank, flags):\n#     torch.set_default_tensor_type('torch.FloatTensor')\n    _test()\n\nFLAGS={}\nif not TRAIN_MODE:\n    xmp.spawn(_mp_fn, args=(FLAGS,), nprocs=1, start_method='fork')","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}