{"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_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## Late submission study\n### Study targets \n* check the full train dataset benefit, as [this competition's winners noticed](https://www.kaggle.com/c/lyft-motion-prediction-autonomous-vehicles/discussion/199657).\n* check [CI and other tips for code quality](https://www.kaggle.com/c/lyft-motion-prediction-autonomous-vehicles/discussion/199506) by [iglovikov](https://www.kaggle.com/iglovikov).\n* to remove manual code writings, use pytorch-lighitning framework as possible.\n\n### Study summary\n* full train dataset benefit: using only 2 history frames but I got 11.376 with 1 epoch train_full.zarr training, \n  it's competitive to the gold medal zone score 11.325.\n* CI and other tips implementations can be checked at [my github](https://github.com/Fkaneko/kaggle-lyft-motion-pred) for this lyft motion prediction.\n  All tips are really nice, especially code fomatter, [black](https://github.com/psf/black) and [pre-commit](https://pre-commit.com/) are impressive. \n* using `pl.LightningDataModule` and `pl.LightningModule`, pytorch-lightinig can cover data loading to submission.","metadata":{}},{"cell_type":"markdown","source":"### Imports and general settings","metadata":{}},{"cell_type":"code","source":"import argparse\nimport os\nimport random\nimport sys\nfrom typing import Tuple\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pytorch_lightning as pl\nimport torch\nimport l5kit # add at kaggle env  \nfrom l5kit.configs import load_config_data\nfrom l5kit.data import ChunkedDataset, LocalDataManager\nfrom l5kit.dataset import AgentDataset\nfrom l5kit.evaluation import compute_metrics_csv, write_pred_csv\nfrom l5kit.evaluation.metrics import neg_multi_log_likelihood, time_displace\nfrom l5kit.geometry import transform_points\nfrom l5kit.rasterization import build_rasterizer\nfrom l5kit.visualization import TARGET_POINTS_COLOR, draw_trajectory\nfrom torch.utils.data import DataLoader\n\nimport lyft_loss\nimport lyft_models\nimport lyft_utils","metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","execution":{"iopub.status.busy":"2022-05-17T01:43:01.976677Z","iopub.execute_input":"2022-05-17T01:43:01.97702Z","iopub.status.idle":"2022-05-17T01:43:17.31275Z","shell.execute_reply.started":"2022-05-17T01:43:01.976988Z","shell.execute_reply":"2022-05-17T01:43:17.311689Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"ALL_DATA_SIZE = 198474478\nVAL_INTERVAL_SAMPLES = 250000\n\n# configuration file for raterization\nCFG_PATH = \"../input/lyft-mpred-seresnext26-pretrained/agent_motion_config.yaml\"\n# output path for test mode\nCSV_PATH = \"./submission.csv\"\n\n# for using the same sampling as test dataset agents,\n# these two FRAME settings are requried.\n# minimum number of frames an agents must have in the past to be picked\nMIN_FRAME_HISTORY = 0\n# minimum number of frames an agents must have in the future to be picked\nMIN_FRAME_FUTURE = 10\nVAL_SELECTED_FRAME = (99,)\n\n# set random seeds\nSEED = 42\nos.environ[\"PYTHONHASHSEED\"] = str(SEED)\n# set Python random seed\nrandom.seed(SEED)\n# set NumPy random seed\nnp.random.seed(SEED)","metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","execution":{"iopub.status.busy":"2022-05-17T01:43:17.315758Z","iopub.execute_input":"2022-05-17T01:43:17.316359Z","iopub.status.idle":"2022-05-17T01:43:17.323555Z","shell.execute_reply.started":"2022-05-17T01:43:17.316316Z","shell.execute_reply":"2022-05-17T01:43:17.322126Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Run configuration","metadata":{}},{"cell_type":"code","source":"cfg = load_config_data(CFG_PATH)\n\nparser = argparse.ArgumentParser(\n    description=\"Run lyft motion prediction learning\",\n    formatter_class=argparse.ArgumentDefaultsHelpFormatter,\n)\nparser.add_argument(\n    \"--l5kit_data_folder\",\n    default=\"/your/dataset/path\",\n    type=str,\n    help=\"root directory path for lyft motion prediction dataset\",\n)\nparser.add_argument(\n    \"--optim_name\",\n    choices=[\"adam\", \"sgd\"],\n    default=\"sgd\",\n    help=\"optimizer name\",\n)\nparser.add_argument(\n    \"--num_modes\",\n    type=int,\n    default=3,\n    help=\"number of the modes on each prediction\",\n)\nparser.add_argument(\"--lr\", default=7.0e-4, type=float, help=\"learning rate\")\nparser.add_argument(\"--batch_size\", type=int, default=220, help=\"batch size\")\nparser.add_argument(\"--epochs\", type=int, default=1, help=\"epochs for training\")\nparser.add_argument(\n    \"--backbone_name\",\n    choices=[\"efficientnet_b1\", \"seresnext26d_32x4d\"],\n    default=\"seresnext26d_32x4d\",\n    help=\"backbone name\",\n)\nparser.add_argument(\n    \"--downsample_train\",\n    action=\"store_true\",\n    help=\"using only 4 frames from each scene, the loss converge is \\\nmuch faster than using all data, but it will get larger loss\",\n)\nparser.add_argument(\"--is_test\", action=\"store_true\", help=\"test mode\")\nparser.add_argument(\n    \"--ckpt_path\",\n    type=str,\n    default=\"./model.pth\",\n    help=\"path for model checkpoint at test mode\",\n)\nparser.add_argument(\n    \"--precision\",\n    default=16,\n    choices=[16, 32],\n    type=int,\n    help=\"float precision at training\",\n)\nparser.add_argument(\n    \"--visible_gpus\",\n    type=str,\n    default=\"0\",\n    help=\"Select gpu ids with comma separated format\",\n)\nparser.add_argument(\n    \"--find_lr\",\n    action=\"store_true\",\n    help=\"find lr with fast ai implementation\",\n)\nparser.add_argument(\n    \"--num_workers\",\n    default=\"16\",\n    type=int,\n    help=\"number of cpus for DataLoader\",\n)\nparser.add_argument(\"--is_debug\", action=\"store_true\", help=\"debug mode\")\n\n# args = parser.parse_args()\nargs = parser.parse_args([\n    \"--l5kit_data_folder\",\n    \"../input/lyft-motion-prediction-autonomous-vehicles\",\n    \"--is_test\",\n    \"--ckpt_path\",\n    \"../input/lyft-mpred-seresnext26-pretrained/epoch-v0.ckpt\",\n   \"--num_workers\",\n    \"4\",\n    \"--batch_size\",\n    \"32\"\n    ])\n\nif args.is_debug:\n    DEBUG = True\n    print(\"\\t ---- DEBUG RUN ---- \")\n    cfg[\"train_data_loader\"][\"key\"] = \"scenes/sample.zarr\"\n    cfg[\"val_data_loader\"][\"key\"] = \"scenes/sample.zarr\"\n    VAL_INTERVAL_SAMPLES = 5000\n    args.batch_size = 16\nelse:\n    DEBUG = False\n    print(\"\\t ---- NORMAL RUN ---- \")\nlyft_utils.print_argparse_arguments(args)\n\nos.environ[\"CUDA_VISIBLE_DEVICES\"] = args.visible_gpus\n","metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","execution":{"iopub.status.busy":"2022-05-17T01:43:17.325022Z","iopub.execute_input":"2022-05-17T01:43:17.325391Z","iopub.status.idle":"2022-05-17T01:43:17.365802Z","shell.execute_reply.started":"2022-05-17T01:43:17.325355Z","shell.execute_reply":"2022-05-17T01:43:17.36511Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Dataloader preparation\n`pl.LightningDataModule` for train, validation and test dataloader.","metadata":{}},{"cell_type":"code","source":"class LyftMpredDatamodule(pl.LightningDataModule):\n    def __init__(\n        self,\n        l5kit_data_folder: str,\n        cfg: dict,\n        batch_size: int = 440,\n        num_workers: int = 16,\n        downsample_train: bool = False,\n        is_test: bool = False,\n        is_debug: bool = False,\n    ) -> None:\n        super().__init__()\n        os.environ[\"L5KIT_DATA_FOLDER\"] = l5kit_data_folder\n        self.cfg = cfg\n        self.batch_size = batch_size\n        self.num_workers = num_workers\n        self.downsample_train = downsample_train\n        self.is_test = is_test\n        self.is_debug = is_debug\n\n    def prepare_data(self):\n        # called only on 1 GPU\n        self.dm = LocalDataManager(None)\n        self.rasterizer = build_rasterizer(cfg, self.dm)\n\n    def setup(self):\n        # called on every GPU\n        if self.is_test:\n            print(\"test mode setup\")\n            self.test_path, test_zarr, self.test_dataset = self.load_zarr_dataset(\n                loader_name=\"test_data_loader\"\n            )\n        else:\n            print(\"train mode setup\")\n            self.train_path, train_zarr, self.train_dataset = self.load_zarr_dataset(\n                loader_name=\"train_data_loader\"\n            )\n            self.val_path, val_zarr, self.val_dataset = self.load_zarr_dataset(\n                loader_name=\"val_data_loader\"\n            )\n            self.plot_dataset(self.train_dataset)\n\n            if self.downsample_train:\n                print(\n                    \"downsampling agents, using only {} frames from each scene\".format(\n                        len(lyft_utils.TRAIN_DSAMPLE_FRAMES)\n                    )\n                )\n                train_agents_list = lyft_utils.downsample_agents(\n                    train_zarr,\n                    self.train_dataset,\n                    selected_frames=lyft_utils.TRAIN_DSAMPLE_FRAMES,\n                )\n                self.train_dataset = torch.utils.data.Subset(\n                    self.train_dataset, train_agents_list\n                )\n            # downsampling the validation dataset same as test dataset or\n            # l5kit.evaluation.create_chopped_dataset\n            val_agents_list = lyft_utils.downsample_agents(\n                val_zarr, self.val_dataset, selected_frames=VAL_SELECTED_FRAME\n            )\n            self.val_dataset = torch.utils.data.Subset(\n                self.val_dataset, val_agents_list\n            )\n\n    def train_dataloader(self):\n        return DataLoader(\n            self.train_dataset,\n            shuffle=True,\n            batch_size=self.batch_size,\n            num_workers=self.num_workers,\n        )\n\n    def val_dataloader(self):\n        return DataLoader(\n            self.val_dataset,\n            shuffle=False,\n            batch_size=self.batch_size,\n            num_workers=self.num_workers,\n        )\n\n    def test_dataloader(self):\n        return DataLoader(\n            self.test_dataset,\n            shuffle=False,\n            batch_size=self.batch_size,\n            num_workers=self.num_workers,\n        )\n\n    def load_zarr_dataset(\n        self, loader_name: str = \"train_data_loder\"\n    ) -> Tuple[str, ChunkedDataset, AgentDataset]:\n\n        zarr_path = self.dm.require(self.cfg[loader_name][\"key\"])\n        print(\"load zarr data:\", zarr_path)\n        zarr_dataset = ChunkedDataset(zarr_path).open()\n        if loader_name == \"test_data_loader\":\n            mask_path = os.path.join(os.path.dirname(zarr_path), \"mask.npz\")\n            agents_mask = np.load(mask_path)[\"arr_0\"]\n            agent_dataset = AgentDataset(\n                self.cfg, zarr_dataset, self.rasterizer, agents_mask=agents_mask\n            )\n        else:\n            agent_dataset = AgentDataset(\n                self.cfg,\n                zarr_dataset,\n                self.rasterizer,\n                min_frame_history=MIN_FRAME_HISTORY,\n                min_frame_future=MIN_FRAME_FUTURE,\n            )\n        print(zarr_dataset)\n        return zarr_path, zarr_dataset, agent_dataset\n\n    def plot_dataset(self, agent_dataset: AgentDataset, plot_num: int = 10) -> None:\n        print(\"Ploting dataset\")\n        ind = np.random.randint(0, len(agent_dataset), size=plot_num)\n        for i in range(plot_num):\n            data = agent_dataset[ind[i]]\n            im = data[\"image\"].transpose(1, 2, 0)\n            im = agent_dataset.rasterizer.to_rgb(im)\n            target_positions_pixels = transform_points(\n                data[\"target_positions\"], data[\"raster_from_agent\"]\n            )\n            draw_trajectory(\n                im,\n                target_positions_pixels,\n                TARGET_POINTS_COLOR,\n                yaws=data[\"target_yaws\"],\n            )\n            plt.imshow(im[::-1])\n            if self.is_debug:\n                plt.show()","metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","execution":{"iopub.status.busy":"2022-05-17T01:43:17.368677Z","iopub.execute_input":"2022-05-17T01:43:17.368933Z","iopub.status.idle":"2022-05-17T01:43:17.396491Z","shell.execute_reply.started":"2022-05-17T01:43:17.368907Z","shell.execute_reply":"2022-05-17T01:43:17.395312Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Train, validation and test steps\n with `pl.LightningModule`.","metadata":{}},{"cell_type":"code","source":"class LitModel(pl.LightningModule):\n    def __init__(\n        self,\n        cfg: dict,\n        num_modes: int = 3,\n        ba_size: int = 128,\n        lr: float = 3.0e-4,\n        backbone_name: str = \"efficientnet_b3\",\n        epochs: int = 1,\n        total_steps: int = 100,\n        data_size: int = ALL_DATA_SIZE,\n        optim_name: str = \"adam\",\n    ) -> None:\n        super().__init__()\n        self.save_hyperparameters(\n            \"lr\",\n            \"backbone_name\",\n            \"num_modes\",\n            \"ba_size\",\n            \"epochs\",\n            \"optim_name\",\n            \"data_size\",\n            \"total_steps\",\n        )\n        self.model = lyft_models.LyftMultiModel(\n            cfg, num_modes=num_modes, backbone_name=backbone_name\n        )\n        self.test_keys = (\"world_from_agent\", \"centroid\", \"timestamp\", \"track_id\")\n\n    def forward(self, x):\n        x = self.model(x)\n        return x\n\n    def training_step(self, batch, batch_idx):\n        inputs = batch[\"image\"]\n        target_availabilities = batch[\"target_availabilities\"].unsqueeze(-1)\n        targets = batch[\"target_positions\"]\n\n        outputs, confidences = self.model(inputs)\n        loss = lyft_loss.pytorch_neg_multi_log_likelihood_batch(\n            targets,\n            outputs,\n            confidences.squeeze(),\n            target_availabilities.squeeze(),\n        )\n        self.log(\n            \"train_epoch_loss\",\n            loss,\n            prog_bar=False,\n            on_epoch=True,\n            on_step=False,\n        )\n        self.log(\"train_loss\", loss)\n        return loss\n\n    def validation_step(self, batch, batch_idx):\n        inputs = batch[\"image\"]\n        target_availabilities = batch[\"target_availabilities\"].unsqueeze(-1)\n        targets = batch[\"target_positions\"]\n\n        outputs, confidences = self.model(inputs)\n        loss = lyft_loss.pytorch_neg_multi_log_likelihood_batch(\n            targets,\n            outputs,\n            confidences.squeeze(),\n            target_availabilities.squeeze(),\n        )\n        self.log(\"val_loss\", loss)\n        return loss\n\n    def test_step(self, batch, batch_idx):\n        inputs = batch[\"image\"]\n        outputs, confidences = self.model(inputs)\n        test_batch = {key_: batch[key_] for key_ in self.test_keys}\n\n        return outputs, confidences, test_batch\n\n    def test_epoch_end(self, outputs):\n        \"\"\"from https://www.kaggle.com/pestipeti/pytorch-baseline-inference\"\"\"\n        pred_coords_list = []\n        confidences_list = []\n        timestamps_list = []\n        track_id_list = []\n\n        # convert into world coordinates and compute offsets\n        for outputs, confidences, batch in outputs:\n            outputs = outputs.cpu().numpy()\n\n            world_from_agents = batch[\"world_from_agent\"].cpu().numpy()\n            centroids = batch[\"centroid\"].cpu().numpy()\n            for idx in range(len(outputs)):\n                for mode in range(3):\n                    outputs[idx, mode, :, :] = (\n                        transform_points(\n                            outputs[idx, mode, :, :], world_from_agents[idx]\n                        )\n                        - centroids[idx][:2]\n                    )\n            pred_coords_list.append(outputs)\n\n            confidences_list.append(confidences)\n            timestamps_list.append(batch[\"timestamp\"])\n            track_id_list.append(batch[\"track_id\"])\n\n        coords = np.concatenate(pred_coords_list)\n        confs = torch.cat(confidences_list).cpu().numpy()\n        track_ids = torch.cat(track_id_list).cpu().numpy()\n        timestamps = torch.cat(timestamps_list).cpu().numpy()\n\n        write_pred_csv(\n            CSV_PATH,\n            timestamps=timestamps,\n            track_ids=track_ids,\n            coords=coords,\n            confs=confs,\n        )\n        print(f\"Saved to {CSV_PATH}\")\n\n    def configure_optimizers(self):\n        if self.hparams.optim_name == \"sgd\":\n            optimizer = torch.optim.SGD(\n                self.parameters(),\n                lr=self.hparams.lr,\n                momentum=0.9,\n                weight_decay=4e-5,\n            )\n        elif self.hparams.optim_name == \"adam\":\n            optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams.lr)\n        else:\n            raise NotImplementedError\n        scheduler = torch.optim.lr_scheduler.OneCycleLR(\n            optimizer, max_lr=self.hparams.lr, total_steps=self.hparams.total_steps\n        )\n        return [optimizer], [scheduler]","metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","execution":{"iopub.status.busy":"2022-05-17T01:43:17.398916Z","iopub.execute_input":"2022-05-17T01:43:17.399253Z","iopub.status.idle":"2022-05-17T01:43:17.429249Z","shell.execute_reply.started":"2022-05-17T01:43:17.399226Z","shell.execute_reply":"2022-05-17T01:43:17.428376Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Instantiating DataModule","metadata":{}},{"cell_type":"code","source":"# ===== Configure LYFT dataset\n# mypy error due to pl.DataModule.transfer_batch_to_device\nmpred_dm = LyftMpredDatamodule(  # type: ignore[abstract]\n    args.l5kit_data_folder,\n    cfg,\n    batch_size=args.batch_size,\n    num_workers=args.num_workers,\n    downsample_train=args.downsample_train,\n    is_test=args.is_test,\n    is_debug=args.is_debug,\n)\nmpred_dm.prepare_data()\nmpred_dm.setup()","metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","execution":{"iopub.status.busy":"2022-05-17T01:43:17.431122Z","iopub.execute_input":"2022-05-17T01:43:17.431857Z","iopub.status.idle":"2022-05-17T01:43:33.748525Z","shell.execute_reply.started":"2022-05-17T01:43:17.431818Z","shell.execute_reply":"2022-05-17T01:43:33.747766Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Running test ","metadata":{}},{"cell_type":"code","source":"if args.is_test:\n    print(\"\\t\\t ==== TEST MODE ====\")\n    print(\"load from: \", args.ckpt_path)\n    model = LitModel.load_from_checkpoint(args.ckpt_path, cfg=cfg)\n    trainer = pl.Trainer(gpus=len(args.visible_gpus.split(\",\")))\n    trainer.test(model, datamodule=mpred_dm)","metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","execution":{"iopub.status.busy":"2022-05-17T01:43:33.751455Z","iopub.execute_input":"2022-05-17T01:43:33.751713Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Running train","metadata":{}},{"cell_type":"code","source":"if not args.is_test:\n    print(\"\\t\\t ==== TRAIN MODE ====\")\n    print(\n        \"training samples: {}, valid samples: {}\".format(\n            len(mpred_dm.train_dataset), len(mpred_dm.val_dataset)\n        )\n    )\n    total_steps = args.epochs * len(mpred_dm.train_dataset) // args.batch_size\n    val_check_interval = VAL_INTERVAL_SAMPLES // args.batch_size\n\n    model = LitModel(\n        cfg,\n        lr=args.lr,\n        backbone_name=args.backbone_name,\n        num_modes=args.num_modes,\n        optim_name=args.optim_name,\n        ba_size=args.batch_size,\n        epochs=args.epochs,\n        data_size=len(mpred_dm.train_dataset),\n        total_steps=total_steps,\n    )\n\n    checkpoint_callback = pl.callbacks.ModelCheckpoint(\n        monitor=\"val_loss\",\n        save_last=True,\n        mode=\"min\",\n        verbose=True,\n    )\n    pl.trainer.seed_everything(seed=SEED)\n    trainer = pl.Trainer(\n        gpus=len(args.visible_gpus.split(\",\")),\n        max_steps=total_steps,\n        val_check_interval=val_check_interval,\n        precision=args.precision,\n        benchmark=True,\n        deterministic=False,\n        checkpoint_callback=checkpoint_callback,\n    )\n    # Run Training\n    trainer.fit(model, datamodule=mpred_dm)","metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Also you can check original code at [my github](https://github.com/Fkaneko/kaggle-lyft-motion-pred).\nTo run the code at kaggle I changed some parts of my original code. ","metadata":{}}]}