Files
caption-test/src/main.py

74 lines
2.1 KiB
Python

"""Entry point: load configs and run training.
Usage: (commands for each preset)
python -m src.main gtauav_balanced
python -m src.main gtauav_balanced_asym
python -m src.main gtauav_balanced_stripnet
python -m src.main gtauav_balanced_stripnet_unfrozen
python -m src.main gtauav_baseline
python -m src.main gtauav_baseline_asym
python -m src.main gtauav_baseline_stripnet
python -m src.main gtauav_baseline_stripnet_unfrozen
python -m src.main gtauav_image_heavy
python -m src.main gtauav_gtauav_text_heavy
"""
from __future__ import annotations
import logging
import sys
import coloredlogs
from src.conf.config_loader import load_all_configs
from src.training.trainer_new import Trainer
from src.utils.path_utils import get_proj_dir
logger = logging.getLogger("caption_test")
def main() -> None:
coloredlogs.install(
level="INFO",
logger=logger,
fmt="%(asctime)s %(name)s %(levelname)s %(message)s",
)
if len(sys.argv) != 2:
raise SystemExit(
"Usage: python -m src.main <preset_name>\n"
"Example: python -m src.main gtauav_balanced\n"
" available presets are subdirectories under in/config_files/",
)
preset_name = sys.argv[1]
proj_dir = get_proj_dir()
path2cfg = f"{proj_dir}in/config_files/" # per REQUIREMENTS_GIN_STYLE.md §5
# -------------------------------------------------------
''' ONLY FOR DEBUG with launch.json config:
"args": ["main gtauav_balanced"] -> so need to extract
preset name "gtauav_balanced"
'''
preset_name = preset_name.split(' ')[1]
# -------------------------------------------------------
configs = load_all_configs(path2cfg, preset_name)
trainer = Trainer(
pipeline_cfg=configs["pipeline"],
hardware_cfg=configs["hardware"],
training_cfg=configs["training"],
tracking_cfg=configs["tracking"],
models_common_cfg=configs["models_common"],
models_cfg=configs["models"],
)
trainer.train()
if __name__ == "__main__":
main()