{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Lyft: Vision Transformer Training"},{"metadata":{},"cell_type":"markdown","source":"This is a kernel on the novel **vision transformer**, and uses the implementation of the model in PyTorch created by **@rwightman** in the python package timm. To install it, we will have to get it straight from the repo (as of now). \n\nFor training abstraction, huge credit to `corochann` and his useful notebook demonstrating how to use PyTorch Ignite and Pytorch-PFN Extras in order to greatly simplify your training routine. So the first question the readers of this kernel might have is, what is a vision transformer?\n\n<h3 style=\"color:red\">The vision transformer explained</h3>\n\n<img src=\"https://github.com/lucidrains/vit-pytorch/raw/main/vit.png\"></img>\n\nMost of you might be familiar with the idea of the NLP transformer (i.e BERT, RoBERTa, XLNET etc.) and what the vision transformer does is apply the fundamental architecture of the NLP transformer (self-attention etc.) within a single encoder of a standard transformer."},{"metadata":{},"cell_type":"markdown","source":"Now for the training, and the fundamental architecture."},{"metadata":{},"cell_type":"markdown","source":"# Setup"},{"metadata":{},"cell_type":"markdown","source":"This covers the libraries, the installation and the setup of the data pipeline. Most of you would already be familiar with this,"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-input":true,"_kg_hide-output":true},"cell_type":"code","source":"!pip install git+https://github.com/rwightman/pytorch-image-models.git\n!pip install --target=/kaggle/working pymap3d==2.1.0 -q\n!pip install --target=/kaggle/working strictyaml -q\n!pip install --target=/kaggle/working protobuf==3.12.2 -q\n!pip install --target=/kaggle/working transforms3d -q\n!pip install --target=/kaggle/working zarr -q\n!pip install --target=/kaggle/working ptable -q\n!pip install --no-dependencies --target=/kaggle/working l5kit==1.1.0 --upgrade -q\n!pip install pytorch-pfn-extras==0.3.1","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"from timm.models import vision_transformer\nimport torch\nimport l5kit, os\nimport torch.nn as nn\nimport numpy\nimport warnings;warnings.filterwarnings(\"ignore\")\nfrom l5kit.rasterization import build_rasterizer\nfrom l5kit.configs import load_config_data\nfrom l5kit.visualization import draw_trajectory, TARGET_POINTS_COLOR\nfrom tqdm import tqdm\nfrom l5kit.geometry import transform_points\nfrom collections import Counter\nfrom l5kit.data import PERCEPTION_LABELS\nfrom prettytable import PrettyTable\nfrom ignite.engine import Events, Engine, create_supervised_trainer\nimport torch.optim as optim\nimport pytorch_pfn_extras.training.extensions as E\nfrom pytorch_pfn_extras.training import IgniteExtensionsManager\nfrom l5kit.evaluation.metrics import neg_multi_log_likelihood\ncfg = load_config_data(\"../input/lyft-config-files/agent_motion_config.yaml\")\nos.environ[\"L5KIT_DATA_FOLDER\"] = \"../input/lyft-motion-prediction-autonomous-vehicles\"\nmodel = vision_transformer.vit_small_resnet50d_s3_224(pretrained=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from l5kit.data import ChunkedDataset, LocalDataManager\nfrom l5kit.dataset import EgoDataset, AgentDataset\ndm = LocalDataManager()\ntrain_cfg = cfg[\"train_data_loader\"]\nrasterizer = build_rasterizer(cfg, dm)\ntrain_zarr = ChunkedDataset(dm.require(train_cfg[\"key\"])).open()\ntrain_dataset = AgentDataset(cfg, train_zarr, rasterizer)\ntrain_dataloader = torch.utils.data.DataLoader(train_dataset,\n                              shuffle=train_cfg[\"shuffle\"],\n                              batch_size=train_cfg[\"batch_size\"],\n                              num_workers=train_cfg[\"num_workers\"])\n\nval_cfg = cfg[\"val_data_loader\"]\nval_zarr = ChunkedDataset(dm.require(val_cfg[\"key\"])).open()\nval_dataset = AgentDataset(cfg, val_zarr, rasterizer)\nval_dataset = torch.utils.data.Subset(val_dataset, range(0, 4000))\nval_dataloader = torch.utils.data.DataLoader(val_dataset,\n                              shuffle=train_cfg[\"shuffle\"],\n                              batch_size=train_cfg[\"batch_size\"],\n                              num_workers=train_cfg[\"num_workers\"])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model"},{"metadata":{},"cell_type":"markdown","source":"The model is a vision transformer with a simple linear layer at the end."},{"metadata":{"trusted":true},"cell_type":"code","source":"class LyftVIT(nn.Module):\n    \n    def __init__(self, vit: nn.Module):\n        super(LyftVIT, self).__init__()\n        num_history_channels = (cfg[\"model_params\"][\"history_num_frames\"] + 1) * 2\n        num_in_channels = 3 + num_history_channels\n        self.vit = vit\n        num_targets = 2 * cfg[\"model_params\"][\"future_num_frames\"]\n        self.future_len = cfg[\"model_params\"][\"future_num_frames\"]\n        self.vit.patch_embed.backbone.conv1[0] = nn.Conv2d(\n            num_in_channels,\n            32,\n            kernel_size=self.vit.patch_embed.backbone.conv1[0].kernel_size,\n            stride=self.vit.patch_embed.backbone.conv1[0].stride,\n            padding=self.vit.patch_embed.backbone.conv1[0].padding,\n            bias=False,\n        )\n        \n        \n        self.num_preds = num_targets * 3\n        self.num_modes = 3\n        \n        self.logit = nn.Linear(1000, out_features=self.num_preds + self.num_modes)\n        \n    def forward(self, x):\n        x = self.vit(x)\n        x = torch.flatten(x, 1)\n        x = self.logit(x)\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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = LyftVIT(model)\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\nmodel.to(device)\noptimizer = optim.Adam(model.parameters(), lr=1e-3)\ncriterion = nn.MSELoss(reduction=\"none\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training"},{"metadata":{"_kg_hide-input":true,"trusted":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":{"trusted":true},"cell_type":"code","source":"def train_step(engine, batch):\n    model.train()\n    optimizer.zero_grad()\n    x, y = batch[\"image\"].to(device), batch[\"target_positions\"].to(device)\n    avails = batch[\"target_availabilities\"].unsqueeze(-1).to(device)\n    y_pred, conf = model(x)\n    loss = pytorch_neg_multi_log_likelihood_batch(y, y_pred, conf, avails[:, :, 0])\n    loss.backward()\n    optimizer.step()\n    return loss.item()\n\ntrainer = Engine(train_step)\n\ndef validation_step(image, target_positions, **kwargs):\n    x, y = image.to(device), target_positions.to(device)\n    y_pred = model(x)\n    return y_pred, y","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"valid_evaluator = E.Evaluator(\n    val_dataloader,\n    model,\n    progress_bar=False,\n    eval_func=validation_step,\n)\n\nlog_trigger = (1000, \"iteration\")\nlog_report = E.LogReport(trigger=log_trigger)\nextensions = [\n    log_report,  # Save `log` to file\n    valid_evaluator,\n    # E.FailOnNonNumber()  # Stop training when nan is detected.\n    E.ProgressBarNotebook(update_interval=100),  # Show progress bar during training\n    E.PrintReportNotebook(),  # Show \"log\" on jupyter notebook  \n]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"models = {\"main\": model}\noptimizers = {\"main\": optimizer}\nmanager = IgniteExtensionsManager(\n    trainer,\n    models,\n    optimizers,\n    7,\n    extensions=extensions,\n    out_dir=\"../working\",\n)\nmanager.extend(E.snapshot_object(model, \"predictor.pt\"),\n               trigger=(50, \"iteration\")) ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"manager.iteration = 0\nmanager._iters_per_epoch = len(train_dataloader)\ntrainer.run(train_dataloader, max_epochs=7)","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}