{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.7.6","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":19990,"databundleVersionId":1472735,"sourceType":"competition"},{"sourceId":1779841,"sourceType":"datasetVersion","datasetId":1058170},{"sourceId":43970252,"sourceType":"kernelVersion"},{"sourceId":50216082,"sourceType":"kernelVersion"},{"sourceId":50258316,"sourceType":"kernelVersion"},{"sourceId":50258373,"sourceType":"kernelVersion"}],"dockerImageVersionId":30043,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"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":"pip install timm==0.4.12","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-19T08:57:31.579010Z","iopub.execute_input":"2025-02-19T08:57:31.579257Z","iopub.status.idle":"2025-02-19T08:57:40.361190Z","shell.execute_reply.started":"2025-02-19T08:57:31.579232Z","shell.execute_reply":"2025-02-19T08:57:40.360059Z"}},"outputs":[],"execution_count":null},{"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","trusted":true,"execution":{"iopub.status.busy":"2025-02-19T08:57:40.364158Z","iopub.execute_input":"2025-02-19T08:57:40.364531Z","iopub.status.idle":"2025-02-19T08:57:48.175079Z","shell.execute_reply.started":"2025-02-19T08:57:40.364497Z","shell.execute_reply":"2025-02-19T08:57:48.173954Z"}},"outputs":[],"execution_count":null},{"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","trusted":true,"execution":{"iopub.status.busy":"2025-02-19T08:57:48.176784Z","iopub.execute_input":"2025-02-19T08:57:48.177231Z","iopub.status.idle":"2025-02-19T08:57:48.183453Z","shell.execute_reply.started":"2025-02-19T08:57:48.177184Z","shell.execute_reply":"2025-02-19T08:57:48.182573Z"}},"outputs":[],"execution_count":null},{"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=\"efficientnet_b1\",\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","trusted":true,"execution":{"iopub.status.busy":"2025-02-19T08:57:48.184581Z","iopub.execute_input":"2025-02-19T08:57:48.184996Z","iopub.status.idle":"2025-02-19T08:57:48.236713Z","shell.execute_reply.started":"2025-02-19T08:57:48.184949Z","shell.execute_reply":"2025-02-19T08:57:48.235672Z"}},"outputs":[],"execution_count":null},{"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 = True,\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","trusted":true,"execution":{"iopub.status.busy":"2025-02-19T08:57:48.238356Z","iopub.execute_input":"2025-02-19T08:57:48.238702Z","iopub.status.idle":"2025-02-19T08:57:48.275367Z","shell.execute_reply.started":"2025-02-19T08:57:48.238668Z","shell.execute_reply":"2025-02-19T08:57:48.274099Z"}},"outputs":[],"execution_count":null},{"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_b1\",\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","trusted":true,"execution":{"iopub.status.busy":"2025-02-19T08:57:48.277417Z","iopub.execute_input":"2025-02-19T08:57:48.277751Z","iopub.status.idle":"2025-02-19T08:57:48.292687Z","shell.execute_reply.started":"2025-02-19T08:57:48.277715Z","shell.execute_reply":"2025-02-19T08:57:48.291745Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport torch\nimport pytorch_lightning as pl\nimport matplotlib.pyplot as plt\nfrom l5kit.geometry import transform_points\n\nclass 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_b1\",\n        epochs: int = 1,\n        total_steps: int = 100,\n        data_size: int = 198474478,\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\", \"target_positions\")\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(\"train_loss\", loss, on_step=False, on_epoch=True)\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, on_step=False, on_epoch=True)\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 {\n            \"outputs\": outputs,\n            \"confidences\": confidences,\n            \"test_batch\": test_batch\n        }\n\n    def test_epoch_end(self, outputs):\n        # Collect all outputs and test batches\n        all_outputs = []\n        all_confidences = []\n        all_world_from_agent = []\n        all_centroids = []\n        all_targets = []\n\n        for output in outputs:\n            all_outputs.append(output[\"outputs\"])\n            all_confidences.append(output[\"confidences\"])\n            test_batch = output[\"test_batch\"]\n            all_world_from_agent.append(test_batch[\"world_from_agent\"])\n            all_centroids.append(test_batch[\"centroid\"])\n            all_targets.append(test_batch[\"target_positions\"])\n\n        # Concatenate all tensors\n        all_outputs = torch.cat(all_outputs).cpu().numpy()\n        all_confidences = torch.cat(all_confidences).cpu().numpy()\n        all_world_from_agent = torch.cat(all_world_from_agent).cpu().numpy()\n        all_centroids = torch.cat(all_centroids).cpu().numpy()\n        all_targets = torch.cat(all_targets).cpu().numpy()\n\n        # Visualize predictions\n        self.visualize_predictions(all_outputs, all_confidences, all_world_from_agent, all_centroids, all_targets)\n\n    def visualize_predictions(self, preds, confidences, world_from_agents, centroids, targets):\n        num_samples = 5  # Number of samples to visualize\n        sample_indices = np.random.choice(len(preds), num_samples, replace=False)\n\n        for idx in sample_indices:\n            pred = preds[idx]  # Shape (num_modes, future_len, 2)\n            conf = confidences[idx]  # Shape (num_modes,)\n            target = targets[idx]  # Shape (future_len, 2)\n            wfa = world_from_agents[idx]\n            centroid = centroids[idx]\n\n            # Transform predictions to world coordinates\n            pred_world = [\n                transform_points(mode, wfa) - centroid[:2]  # Subtract centroid to get world coordinates\n                for mode in pred\n            ]\n\n            # Plotting\n            plt.figure(figsize=(8, 8))\n            plt.title(f\"Sample {idx} - Prediction vs Target\")\n            plt.plot(target[:, 0], target[:, 1], marker='o', label=\"Target\", color='blue')\n\n            # Plot each prediction mode\n            for mode_idx in range(len(pred_world)):\n                plt.plot(\n                    pred_world[mode_idx][:, 0],\n                    pred_world[mode_idx][:, 1],\n                    marker='x',\n                    linestyle='--',\n                    label=f\"Mode {mode_idx} (Conf: {conf[mode_idx]:.2f})\",\n                    color=plt.cm.RdBu(mode_idx / len(pred_world))\n                )\n\n            plt.legend()\n            plt.xlabel(\"X (World)\")\n            plt.ylabel(\"Y (World)\")\n            plt.grid()\n            plt.show()\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":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-19T08:57:48.294260Z","iopub.execute_input":"2025-02-19T08:57:48.294507Z","iopub.status.idle":"2025-02-19T08:57:48.330150Z","shell.execute_reply.started":"2025-02-19T08:57:48.294482Z","shell.execute_reply":"2025-02-19T08:57:48.329156Z"}},"outputs":[],"execution_count":null},{"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","trusted":true,"execution":{"iopub.status.busy":"2025-02-19T08:57:48.332198Z","iopub.execute_input":"2025-02-19T08:57:48.332485Z","iopub.status.idle":"2025-02-19T08:58:10.263234Z","shell.execute_reply.started":"2025-02-19T08:57:48.332456Z","shell.execute_reply":"2025-02-19T08:58:10.262402Z"}},"outputs":[],"execution_count":null},{"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","trusted":true,"execution":{"iopub.status.busy":"2025-02-19T08:58:10.267005Z","iopub.execute_input":"2025-02-19T08:58:10.267252Z","iopub.status.idle":"2025-02-19T08:58:10.270872Z","shell.execute_reply.started":"2025-02-19T08:58:10.267225Z","shell.execute_reply":"2025-02-19T08:58:10.270098Z"}},"outputs":[],"execution_count":null},{"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\n    # # 运行测试以获取预测结果\n    trainer.test(model, datamodule=mpred_dm)\n\n    # # 绘制训练损失曲线（如果需要）\n    # train_loss = [metric[\"train_loss\"] for metric in trainer.logged_metrics.values() if \"train_loss\" in metric]\n    # val_loss = [metric[\"val_loss\"] for metric in trainer.logged_metrics.values() if \"val_loss\" in metric]\n\n    # plt.figure(figsize=(10, 6))\n    # plt.plot(train_loss, label=\"Training Loss\", color=\"blue\", marker=\"o\")\n    # plt.plot(val_loss, label=\"Validation Loss\", color=\"red\", marker=\"o\")\n    # plt.title(\"Training and Validation Loss\")\n    # plt.xlabel(\"Epoch\")\n    # plt.ylabel(\"Loss\")\n    # plt.legend()\n    # plt.grid()\n    # plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-19T08:58:30.527525Z","iopub.execute_input":"2025-02-19T08:58:30.527875Z","iopub.status.idle":"2025-02-19T09:48:32.900333Z","shell.execute_reply.started":"2025-02-19T08:58:30.527843Z","shell.execute_reply":"2025-02-19T09:48:32.899408Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_loss = [metric[\"train_loss\"] for metric in trainer.logged_metrics.values() if \"train_loss\" in metric]\nval_loss = [metric[\"val_loss\"] for metric in trainer.logged_metrics.values() if \"val_loss\" in metric]\n\nplt.figure(figsize=(10, 6))\nplt.plot(train_loss, label=\"Training Loss\", color=\"blue\", marker=\"o\")\nplt.plot(val_loss, label=\"Validation Loss\", color=\"red\", marker=\"o\")\nplt.title(\"Training and Validation Loss\")\nplt.xlabel(\"Epoch\")\nplt.ylabel(\"Loss\")\nplt.legend()\nplt.grid()\nplt.show()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"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":{"iopub.status.busy":"2025-02-19T08:58:14.206315Z","iopub.execute_input":"2025-02-19T08:58:14.206633Z","iopub.status.idle":"2025-02-19T08:58:14.217935Z","shell.execute_reply.started":"2025-02-19T08:58:14.206602Z","shell.execute_reply":"2025-02-19T08:58:14.216845Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import matplotlib.pyplot as plt","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-18T05:49:00.015117Z","iopub.execute_input":"2025-02-18T05:49:00.015419Z","iopub.status.idle":"2025-02-18T05:49:00.019245Z","shell.execute_reply.started":"2025-02-18T05:49:00.015394Z","shell.execute_reply":"2025-02-18T05:49:00.018294Z"}},"outputs":[],"execution_count":null},{"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":{}},{"cell_type":"code","source":"import matplotlib.pyplot as plt\n\ndef visualize_predictions(self, preds, confidences, test_batches):\n    # 随机选择几个 agent 进行可视化\n    num_samples = 5\n    sample_indices = np.random.choice(len(preds), num_samples, replace=False)\n\n    for idx in sample_indices:\n        pred = preds[idx]  # shape (num_modes, future_len, 2)\n        conf = confidences[idx]  # shape (num_modes,)\n        target = test_batches[\"target_positions\"][idx]  # shape (future_len, 2)\n        world_from_agent = test_batches[\"world_from_agent\"][idx]\n        centroid = test_batches[\"centroid\"][idx]\n\n        # 将预测结果转换为世界坐标\n        pred_world = [\n            transform_points(mode, world_from_agent) - centroid[:2]\n            for mode in pred\n        ]\n\n        # 绘制图像\n        plt.figure(figsize=(8, 8))\n        plt.title(f\"Sample {idx} - Prediction and Target\")\n\n        # 绘制真实轨迹\n        plt.plot(target[:, 0], target[:, 1], marker=\"o\", label=\"Target\", color=\"blue\")\n\n        # 绘制预测轨迹\n        for mode_idx in range(pred.shape[0]):\n            plt.plot(\n                pred_world[mode_idx][:, 0],\n                pred_world[mode_idx][:, 1],\n                marker=\"x\",\n                linestyle=\"--\",\n                label=f\"Mode {mode_idx} (Conf: {conf[mode_idx]:.2f})\",\n                color=plt.cm.RdBu(mode_idx / len(conf))\n            )\n\n        plt.legend()\n        plt.xlabel(\"X (world coordinates)\")\n        plt.ylabel(\"Y (world coordinates)\")\n        plt.grid()\n        plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-18T05:51:21.090688Z","iopub.execute_input":"2025-02-18T05:51:21.090993Z","iopub.status.idle":"2025-02-18T05:51:21.100849Z","shell.execute_reply.started":"2025-02-18T05:51:21.090968Z","shell.execute_reply":"2025-02-18T05:51:21.099906Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ... 之前的代码不变 ...\n\nif 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)\n\n    # Run Test to visualize predictions\n    print(\"Running test to visualize predictions...\")\n    trainer.test(model, datamodule=mpred_dm)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-18T05:51:49.597989Z","iopub.execute_input":"2025-02-18T05:51:49.598278Z","iopub.status.idle":"2025-02-18T05:51:49.606445Z","shell.execute_reply.started":"2025-02-18T05:51:49.598253Z","shell.execute_reply":"2025-02-18T05:51:49.605674Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}