|
import time |
|
import torch |
|
import hydra |
|
import pytorch_lightning as pl |
|
from typing import Any |
|
|
|
from hydra.core.config_store import ConfigStore |
|
from omegaconf import OmegaConf |
|
from pytorch_lightning.loggers import WandbLogger |
|
from pytorch_lightning.callbacks import ModelCheckpoint |
|
|
|
from pathlib import Path |
|
from dataclasses import dataclass |
|
|
|
from .module import GenericModule |
|
from .data.module import GenericDataModule |
|
from .callbacks import EvalSaveCallback, ImageLoggerCallback |
|
from .models.schema import ModelConfiguration, DINOConfiguration, ResNetConfiguration |
|
from .data.schema import MIADataConfiguration, KITTIDataConfiguration, NuScenesDataConfiguration |
|
|
|
|
|
@dataclass |
|
class ExperimentConfiguration: |
|
name: str |
|
|
|
@dataclass |
|
class Configuration: |
|
model: ModelConfiguration |
|
experiment: ExperimentConfiguration |
|
data: Any |
|
training: Any |
|
|
|
|
|
cs = ConfigStore.instance() |
|
|
|
|
|
cs.store(name="pretrain", node=Configuration) |
|
cs.store(name="mapper_nuscenes", node=Configuration) |
|
cs.store(name="mapper_kitti", node=Configuration) |
|
|
|
|
|
cs.store(group="schema/data", name="mia", |
|
node=MIADataConfiguration, package="data") |
|
cs.store(group="schema/data", name="kitti", node=KITTIDataConfiguration, package="data") |
|
cs.store(group="schema/data", name="nuscenes", node=NuScenesDataConfiguration, package="data") |
|
|
|
cs.store(group="model/schema/backbone", name="dino", node=DINOConfiguration, package="model.image_encoder.backbone") |
|
cs.store(group="model/schema/backbone", name="resnet", node=ResNetConfiguration, package="model.image_encoder.backbone") |
|
|
|
|
|
@hydra.main(version_base=None, config_path="conf", config_name="pretrain") |
|
def train(cfg: Configuration): |
|
OmegaConf.resolve(cfg) |
|
|
|
dm = GenericDataModule(cfg.data) |
|
|
|
model = GenericModule(cfg) |
|
|
|
exp_name_with_time = cfg.experiment.name + \ |
|
"_" + time.strftime("%Y-%m-%d_%H-%M-%S") |
|
|
|
callbacks: list[pl.Callback] |
|
|
|
if cfg.training.eval: |
|
save_dir = Path(cfg.training.save_dir) |
|
save_dir.mkdir(parents=True, exist_ok=True) |
|
|
|
callbacks = [ |
|
EvalSaveCallback(save_dir=save_dir) |
|
] |
|
|
|
logger = None |
|
else: |
|
callbacks = [ |
|
ImageLoggerCallback(num_classes=cfg.training.num_classes), |
|
ModelCheckpoint( |
|
monitor=cfg.training.checkpointing.monitor, |
|
save_last=cfg.training.checkpointing.save_last, |
|
save_top_k=cfg.training.checkpointing.save_top_k, |
|
) |
|
] |
|
|
|
logger = WandbLogger( |
|
name=exp_name_with_time, |
|
id=exp_name_with_time, |
|
entity="mappred-large", |
|
project="map-pred-full-v3", |
|
) |
|
|
|
logger.watch(model, log="all", log_freq=500) |
|
|
|
if cfg.training.checkpoint is not None: |
|
state_dict = torch.load(cfg.training.checkpoint)['state_dict'] |
|
model.load_state_dict(state_dict, strict=False) |
|
|
|
trainer_args = OmegaConf.to_container(cfg.training.trainer) |
|
trainer_args['callbacks'] = callbacks |
|
trainer_args['logger'] = logger |
|
|
|
trainer = pl.Trainer(**trainer_args) |
|
|
|
if cfg.training.eval: |
|
trainer.test(model, datamodule=dm) |
|
else: |
|
trainer.fit(model, datamodule=dm) |
|
|
|
|
|
if __name__ == "__main__": |
|
pl.seed_everything(42) |
|
torch.set_float32_matmul_precision("high") |
|
|
|
train() |
|
|
