{"cells":[{"metadata":{},"cell_type":"markdown","source":"## Motivation\n* There are some tpu kernels in this competition,but no one try to use trainer( Shams Sultan may try, but it seems to fail. https://www.kaggle.com/shams1/lyft-understanding-the-data-baseline-model)\n* By using trainer, we can easily use tpu. So, I explain it in this kernel.\n* This kernel is firmly based on corochann's kernel.Thanks. **[Lyft: Prediction with multi-mode confidence]**(https://www.kaggle.com/corochann/lyft-prediction-with-multi-mode-confidence)."},{"metadata":{},"cell_type":"markdown","source":"## Install some modules"},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"!pip install torch==1.6.0","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-output":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,"_kg_hide-input":true,"_kg_hide-output":true},"cell_type":"code","source":"!pip install pytorch-lightning==1.0.4\nimport pytorch_lightning as pl","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import pytorch_lightning as pl\npl.__version__","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"import gc\nimport os\nfrom pathlib import Path\nimport random\nimport sys\n\nfrom tqdm.notebook import tqdm\nimport numpy as np\nimport pandas as pd\nimport scipy as sp\n\n\nimport matplotlib.pyplot as plt\nimport seaborn as sns\n\nfrom IPython.core.display import display, HTML\nimport zarr\n\nimport l5kit\nfrom l5kit.data import ChunkedDataset, LocalDataManager\nfrom l5kit.dataset import EgoDataset, AgentDataset\n\nfrom l5kit.rasterization import build_rasterizer\nfrom l5kit.configs import load_config_data\nfrom l5kit.visualization import draw_trajectory, TARGET_POINTS_COLOR\nfrom l5kit.geometry import transform_points\nfrom tqdm import tqdm\nfrom collections import Counter\nfrom l5kit.data import PERCEPTION_LABELS\n\nfrom prettytable import PrettyTable\n\nfrom matplotlib import animation, rc\nfrom IPython.display import HTML\n\n\nrc('animation', html='jshtml')\nprint(\"l5kit version:\", l5kit.__version__)\nimport torch\nfrom pathlib import Path\nfrom math import ceil\nfrom torch import nn, optim\nfrom torch.utils.data import DataLoader\nfrom torch.utils.data.dataset import Subset","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"from typing import Callable\n\nfrom torch.utils.data.dataset import Dataset\n\n\nclass TransformDataset(Dataset):\n    def __init__(self, dataset: Dataset, transform: Callable):\n        self.dataset = dataset\n        self.transform = transform\n\n    def __getitem__(self, index):\n        batch = self.dataset[index]\n        return self.transform(batch)\n\n    def __len__(self):\n        return len(self.dataset)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"# --- 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    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":{},"cell_type":"markdown","source":"## Model\n\nWe have to inherit pl.LightningModule, and to make training_step and configure_optimizers."},{"metadata":{"trusted":true,"_kg_hide-input":false},"cell_type":"code","source":"# --- Model utils ---\nimport torch\nfrom torchvision.models import resnet18\nfrom torch import nn\nfrom typing import Dict\n\n\nclass LyftMultiModel(pl.LightningModule):\n\n    def __init__(self, num_modes=3):\n        super().__init__()\n\n        # TODO: support other than resnet18?\n        backbone = resnet18(pretrained=True, progress=True)\n        self.backbone = backbone\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        self.future_len = cfg[\"model_params\"][\"future_num_frames\"]\n        num_targets = 2 * self.future_len\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.num_preds = num_targets * num_modes\n        self.num_modes = num_modes\n\n        self.logit = nn.Linear(4096, out_features=self.num_preds + num_modes)\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        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        return pred, confidences\n\n    def training_step(self, batch, batch_idx: int): \n        target_availabilities = batch[-1]\n        targets = batch[-2]\n        im = batch[0]\n        pred, confidences = self(im)\n        loss = pytorch_neg_multi_log_likelihood_batch(targets, pred, confidences, target_availabilities)\n        return loss\n\n    def validation_step(self, batch, batch_idx: int):\n        target_availabilities = batch[-1]\n        targets = batch[-2]\n        im = batch[0]\n        pred, confidences = self(im)\n        loss = pytorch_neg_multi_log_likelihood_batch(targets, pred, confidences, target_availabilities)\n        return loss\n#         self.log(\"val_loss\", loss)\n    def configure_optimizers(self):\n        return [torch.optim.Adam(self.parameters(), lr=0.000001)]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"# --- Utils ---\nimport yaml\n\n\ndef save_yaml(filepath, content, width=120):\n    with open(filepath, 'w') as f:\n        yaml.dump(content, f, width=width)\n\n\ndef load_yaml(filepath):\n    with open(filepath, 'r') as f:\n        content = yaml.safe_load(f)\n    return content\n\n\nclass DotDict(dict):\n    \"\"\"dot.notation access to dictionary attributes\n\n    Refer: https://stackoverflow.com/questions/2352181/how-to-use-a-dot-to-access-members-of-dictionary/23689767#23689767\n    \"\"\"  # NOQA\n\n    __getattr__ = dict.get\n    __setattr__ = dict.__setitem__\n    __delattr__ = dict.__delitem__\n\n    ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Configs"},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"# --- Lyft configs ---\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': [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    'train_data_loader': {\n        'key': 'scenes/train.zarr',\n        'batch_size': 12,\n        'shuffle': True,\n        'num_workers': 4\n    },\n\n    'valid_data_loader': {\n        'key': 'scenes/validate.zarr',\n        'batch_size': 32,\n        'shuffle': False,\n        'num_workers': 4\n    },\n\n    'train_params': {\n        'max_num_steps': 10000,\n        # 'eval_every_n_steps': -1\n    }\n}\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"flags_dict = {\n    \"debug\": True,\n    # --- Data configs ---\n    \"l5kit_data_folder\": \"/kaggle/input/lyft-motion-prediction-autonomous-vehicles\",\n    # --- Model configs ---\n    \"pred_mode\": \"multi\",\n    # --- Training configs ---\n    \"device\": \"cuda:0\",\n    \"out_dir\": \"results/multi_train\",\n    \"epoch\": 2,\n    \"snapshot_freq\": 50,\n}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"flags = DotDict(flags_dict)\nout_dir = Path(flags.out_dir)\nos.makedirs(str(out_dir), exist_ok=True)\nprint(f\"flags: {flags_dict}\")\nsave_yaml(out_dir / 'flags.yaml', flags_dict)\nsave_yaml(out_dir / 'cfg.yaml', cfg)\ndebug = flags.debug","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Dataset\nIs is compeletely same as base kernel."},{"metadata":{"trusted":true,"_kg_hide-input":false,"_kg_hide-output":false},"cell_type":"code","source":"# set env variable for data\nos.environ[\"L5KIT_DATA_FOLDER\"] = flags.l5kit_data_folder\ndm = LocalDataManager(None)\n\nprint(\"Load dataset...\")\ntrain_cfg = cfg[\"train_data_loader\"]\nvalid_cfg = cfg[\"valid_data_loader\"]\n\n# Rasterizer\nrasterizer = build_rasterizer(cfg, dm)\n\n# Train dataset/dataloader\ndef transform(batch):\n    return batch[\"image\"], batch[\"target_positions\"], batch[\"target_availabilities\"]\n\n\ntrain_path = \"scenes/sample.zarr\" if debug else train_cfg[\"key\"]\ntrain_zarr = ChunkedDataset(dm.require(train_path)).open()\nprint(\"train_zarr\", type(train_zarr))\ntrain_agent_dataset = AgentDataset(cfg, train_zarr, rasterizer)\ntrain_dataset = TransformDataset(train_agent_dataset, transform)\nif debug:\n    # Only use 1000 dataset for fast check...\n    train_dataset = Subset(train_dataset, np.arange(1000))\ntrain_loader = DataLoader(train_dataset,\n                          shuffle=train_cfg[\"shuffle\"],\n                          batch_size=train_cfg[\"batch_size\"],\n                          num_workers=train_cfg[\"num_workers\"])\nprint(train_agent_dataset)\n\nvalid_path = \"scenes/sample.zarr\" if debug else valid_cfg[\"key\"]\nvalid_zarr = ChunkedDataset(dm.require(valid_path)).open()\nprint(\"valid_zarr\", type(train_zarr))\nvalid_agent_dataset = AgentDataset(cfg, valid_zarr, rasterizer)\nvalid_dataset = TransformDataset(valid_agent_dataset, transform)\nif debug:\n    # Only use 100 dataset for fast check...\n    valid_dataset = Subset(valid_dataset, np.arange(100))\nelse:\n    # Only use 1000 dataset for fast check...\n    valid_dataset = Subset(valid_dataset, np.arange(1000))\nvalid_loader = DataLoader(\n    valid_dataset,\n    shuffle=valid_cfg[\"shuffle\"],\n    batch_size=valid_cfg[\"batch_size\"],\n    num_workers=valid_cfg[\"num_workers\"]\n)\n\nprint(valid_agent_dataset)\nprint(\"# AgentDataset train:\", len(train_agent_dataset), \"#valid\", len(valid_agent_dataset))\nprint(\"# ActualDataset train:\", len(train_dataset), \"#valid\", len(valid_dataset))\n# AgentDataset train: 22496709 #valid 21624612\n# ActualDataset train: 100 #valid 100","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"trainer = pl.Trainer(\n    tpu_cores = 1,\n    max_epochs = 1,\n# #     gpus=1,\n#     max_steps=10,\n    amp_backend='native',\n#     default_root_dir='./',\n    profiler=\"simple\",\n#     early_stop_callback=None\n    limit_train_batches = 10,\n    checkpoint_callback = False,\n#     weights_save_path = \"./a\"\n    logger = False\n    \n    )","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Train\ntrainer can be simple. what we have to get ready for is only model and lodaer. "},{"metadata":{"trusted":true},"cell_type":"code","source":"predictor = LyftMultiModel()\npt_path = \"/kaggle/input/lyft-resnet18-baseline/0918_predictor_full.pt\"\nprint(f\"Loading from {pt_path}\")\npredictor.load_state_dict(torch.load(pt_path,map_location=torch.device('cpu')))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pl.Trainer?","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":""},{"metadata":{"trusted":true,"_kg_hide-input":false},"cell_type":"code","source":"trainer.fit(predictor, train_loader,valid_loader)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# vars(trainer).checkpoint","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}