{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Lyft: Vision Transformer Inference + TTA"},{"metadata":{},"cell_type":"markdown","source":"This is now the inference for the earlier notebook, [Lyft Vision Transformer Training](https://www.kaggle.com/nxrprime/lyft-vision-transformer-training). "},{"metadata":{"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","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Setup"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"from timm.models import vision_transformer\nimport torch\nimport l5kit, os\nimport torch.nn as nn\nimport numpy as np\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\nimport torch.optim as optim\nimport torch.nn.functional as F\nfrom l5kit.evaluation import write_pred_csv\nfrom l5kit.geometry import transform_points\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,"_kg_hide-input":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': 4\n    }\n\n}","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()\ntest_config = cfg[\"test_data_loader\"]\nrasterizer = build_rasterizer(cfg, dm)\ntest_chunked = ChunkedDataset(dm.require(test_config[\"key\"])).open()\ntest_mask = np.load(f\"../input/lyft-motion-prediction-autonomous-vehicles/scenes/mask.npz\")[\"arr_0\"]\ntest_dataset = AgentDataset(cfg, test_chunked, rasterizer, agents_mask=test_mask)\ntest_loader = torch.utils.data.DataLoader(test_dataset,\n                              shuffle=test_config[\"shuffle\"],\n                              batch_size=test_config[\"batch_size\"],\n                              num_workers=test_config[\"num_workers\"])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"class LyftVIT(nn.Module):\n    \n    def __init__(self, vit: nn.Module):\n        super().__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.Conv1d(\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\n    \nmodel = LyftVIT(model)\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\nmodel.to(device)\nprint(\"Model initialized.\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\nmodel.vit.patch_embed.backbone.conv1[0] = nn.Conv2d(\n            5,\n            32,\n            kernel_size=model.vit.patch_embed.backbone.conv1[0].kernel_size,\n            stride=model.vit.patch_embed.backbone.conv1[0].stride,\n            padding=model.vit.patch_embed.backbone.conv1[0].padding,\n            bias=False,\n        )\nmodel.load_state_dict(torch.load('../input/lyft-vision-transformer-training/predictor.pt'))\nmodel.vit.patch_embed.backbone.conv1[0] = nn.Conv2d(\n            25,\n            32,\n            kernel_size=model.vit.patch_embed.backbone.conv1[0].kernel_size,\n            stride=model.vit.patch_embed.backbone.conv1[0].stride,\n            padding=model.vit.patch_embed.backbone.conv1[0].padding,\n            bias=False,\n        ).to(device)\n\n# this is a bit hacky, kinda imperfect.\n# the main issue is to transfer the weights with 25-channel input images","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Inference"},{"metadata":{"trusted":true},"cell_type":"code","source":"model.eval()\n\nfuture_coords_offsets_pd = []\ntimestamps = []\nagent_ids = []\nconfs = []\n\nwith torch.no_grad():\n    dataiter = tqdm(test_loader)\n    \n    for data in dataiter:\n\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, conf = model(inputs)\n        preds = outputs.cpu().numpy()\n        conf = conf.cpu().numpy()\n        world_from_agents = data[\"world_from_agent\"].numpy()\n        centroids = data[\"centroid\"].numpy()\n        coords_offset = []\n        \n        # convert into world coordinates and compute offsets\n        for idx in range(len(preds)):\n            for mode in range(3):\n                preds[idx, mode, :, :] = transform_points(preds[idx, mode, :, :], world_from_agents[idx]) - centroids[idx][:2]\n    \n        future_coords_offsets_pd.append(preds.copy())\n        confs.append(conf.copy())\n        timestamps.append(data[\"timestamp\"].numpy().copy())\n        agent_ids.append(data[\"track_id\"].numpy().copy()) ","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),\n              confs=np.concatenate(confs))","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}