File size: 3,464 Bytes
da2e2ac |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
from typing import Tuple
import hydra
from hydra.utils import instantiate
import logging
from omegaconf import DictConfig
from pathlib import Path
import pytorch_lightning as pl
from torch.utils.data import DataLoader
from det_map.agent_lightning import AgentLightningModuleMap
from det_map.data.datasets.dataset_det import DetDataset
from det_map.utils import collate_fn_pad_lidar
from det_map.data.datasets.dataset import Dataset
from navsim.planning.training.agent_lightning_module import AgentLightningModule
from det_map.data.datasets.dataloader import SceneLoader
from det_map.data.datasets.dataclasses import SceneFilter
from navsim.agents.abstract_agent import AbstractAgent
logger = logging.getLogger(__name__)
CONFIG_PATH = "config/"
CONFIG_NAME = "train_det"
def build_datasets(cfg: DictConfig, agent: AbstractAgent) -> Tuple[Dataset, Dataset]:
train_scene_filter: SceneFilter = instantiate(cfg.scene_filter)
train_scene_filter.log_names = cfg.train_logs
val_scene_filter: SceneFilter = instantiate(cfg.scene_filter)
val_scene_filter.log_names = cfg.val_logs
data_path = Path(cfg.navsim_log_path)
sensor_blobs_path = Path(cfg.sensor_blobs_path)
train_scene_loader = SceneLoader(
sensor_blobs_path=sensor_blobs_path,
data_path=data_path,
scene_filter=train_scene_filter,
sensor_config=agent.get_sensor_config(),
)
val_scene_loader = SceneLoader(
sensor_blobs_path=sensor_blobs_path,
data_path=data_path,
scene_filter=val_scene_filter,
sensor_config=agent.get_sensor_config(),
)
train_data = DetDataset(
scene_loader=train_scene_loader,
feature_builders=agent.get_feature_builders(),
target_builders=agent.get_target_builders(),
pipelines=agent.pipelines,
is_train=True
)
val_data = DetDataset(
scene_loader=val_scene_loader,
feature_builders=agent.get_feature_builders(),
target_builders=agent.get_target_builders(),
pipelines=agent.pipelines,
is_train=False
)
return train_data, val_data
@hydra.main(config_path=CONFIG_PATH, config_name=CONFIG_NAME)
def main(cfg: DictConfig) -> None:
logger.info("Global Seed set to 0")
pl.seed_everything(0, workers=True)
logger.info(f"Path where all results are stored: {cfg.output_dir}")
logger.info("Building Agent")
agent: AbstractAgent = instantiate(cfg.agent)
logger.info("Building Lightning Module")
lightning_module = AgentLightningModuleMap(
agent=agent,
)
logger.info("Building SceneLoader")
train_data, val_data = build_datasets(cfg, agent)
logger.info("Building Datasets")
train_dataloader = DataLoader(train_data, **cfg.dataloader.params, shuffle=True, collate_fn=collate_fn_pad_lidar)
logger.info("Num training samples: %d", len(train_data))
val_dataloader = DataLoader(val_data, **cfg.dataloader.params, shuffle=False, collate_fn=collate_fn_pad_lidar)
logger.info("Num validation samples: %d", len(val_data))
logger.info("Building Trainer")
trainer = pl.Trainer(**cfg.trainer.params)
logger.info("Starting Training")
trainer.fit(
model=lightning_module,
train_dataloaders=train_dataloader,
val_dataloaders=val_dataloader,
)
if __name__ == "__main__":
main()
|