Upload TIGER project for AutoDL

d9115cb verified 9 months ago

4.53 kB

import os
from typing import Any, Dict, Optional, Tuple

import hydra
import rootutils
import torch
from omegaconf import DictConfig

rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)

# ------------------------------------------------------------------------------------ #
# the setup_root above is equivalent to:
# - adding project root dir to PYTHONPATH
#       (so you don't need to force user to install project as a package)
#       (necessary before importing any local modules e.g. `from src import utils`)
# - setting up PROJECT_ROOT environment variable
#       (which is used as a base for paths in "configs/paths/default.yaml")
#       (this way all filepaths are the same no matter where you run the code)
# - loading environment variables from ".env" in root dir
#
# you can remove it if you:
# 1. either install project as a package or move entry files to project root dir
# 2. set `root_dir` to "." in "configs/paths/default.yaml"
#
# more info: https://github.com/ashleve/rootutils
# ------------------------------------------------------------------------------------ #
from src.utils import RankedLogger, extras
from src.utils.custom_hydra_resolvers import *
from src.utils.launcher_utils import pipeline_launcher
from src.utils.restart_job import LocalJobLauncher

command_line_logger = RankedLogger(__name__, rank_zero_only=True)

torch.set_float32_matmul_precision("medium")


def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]:
    """Trains the model. Can additionally evaluate on a testset, using best weights obtained during
    training.

    This method is wrapped in optional @task_wrapper decorator, that controls the behavior during
    failure. Useful for multiruns, saving info about the crash, etc.

    :param cfg: A DictConfig configuration composed by Hydra.
    :return: A tuple with metrics and dict with all instantiated objects.
    """
    # Pipeline launcher initializes the modules needed for the pipeline to run.
    # It also serves as a context manager, so all resources are properly closed after the pipeline is done.
    with pipeline_launcher(cfg) as pipeline_modules:

        if cfg.get("train"):
            command_line_logger.info("Starting training!")
            pipeline_modules.trainer.fit(
                model=pipeline_modules.model,
                datamodule=pipeline_modules.datamodule,
                ckpt_path=cfg.get("ckpt_path"),
            )
        train_metrics = pipeline_modules.trainer.callback_metrics

        if cfg.get("test"):
            command_line_logger.info("Starting testing!")
            ckpt_path = None
            # Check if a checkpoint callback is available and if it has a best model path.
            # Note that if multiple checkpoint callbacks are used, only the first one will be used
            # to determine the best model path for testing.
            checkpoint_callback = getattr(
                pipeline_modules.trainer, "checkpoint_callback", None
            )
            if checkpoint_callback:
                ckpt_path = getattr(checkpoint_callback, "best_model_path", None)
                if ckpt_path == "":
                    ckpt_path = None
            if not ckpt_path:
                command_line_logger.warning(
                    "Best checkpoint not found! Using current weights for testing..."
                )
            pipeline_modules.trainer.test(
                model=pipeline_modules.model,
                datamodule=pipeline_modules.datamodule,
                ckpt_path=ckpt_path,
            )
            command_line_logger.info(f"Best ckpt path: {ckpt_path}")

        test_metrics = pipeline_modules.trainer.callback_metrics

        # merge train and test metrics
        metric_dict = {**train_metrics, **test_metrics}

        command_line_logger.info(f"Metrics: {metric_dict}")


@hydra.main(version_base="1.3", config_path="../configs", config_name="train.yaml")
def main(cfg: DictConfig) -> Optional[float]:
    """Main entry point for training.

    :param cfg: DictConfig configuration composed by Hydra.
    :return: Optional[float] with optimized metric value.
    """
    # apply extra utilities
    # (e.g. ask for tags if none are provided in cfg, print cfg tree, etc.)
    extras(cfg)
    job_launcher = LocalJobLauncher(cfg=cfg)
    job_launcher.launch(function_to_run=train)


if __name__ == "__main__":
    main()

	import os
	from typing import Any, Dict, Optional, Tuple

	import hydra
	import rootutils
	import torch
	from omegaconf import DictConfig

	rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)

	# ------------------------------------------------------------------------------------ #
	# the setup_root above is equivalent to:
	# - adding project root dir to PYTHONPATH
	# (so you don't need to force user to install project as a package)
	# (necessary before importing any local modules e.g. `from src import utils`)
	# - setting up PROJECT_ROOT environment variable
	# (which is used as a base for paths in "configs/paths/default.yaml")
	# (this way all filepaths are the same no matter where you run the code)
	# - loading environment variables from ".env" in root dir
	#
	# you can remove it if you:
	# 1. either install project as a package or move entry files to project root dir
	# 2. set `root_dir` to "." in "configs/paths/default.yaml"
	#
	# more info: https://github.com/ashleve/rootutils
	# ------------------------------------------------------------------------------------ #
	from src.utils import RankedLogger, extras
	from src.utils.custom_hydra_resolvers import *
	from src.utils.launcher_utils import pipeline_launcher
	from src.utils.restart_job import LocalJobLauncher

	command_line_logger = RankedLogger(__name__, rank_zero_only=True)

	torch.set_float32_matmul_precision("medium")


	def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]:
	"""Trains the model. Can additionally evaluate on a testset, using best weights obtained during
	training.

	This method is wrapped in optional @task_wrapper decorator, that controls the behavior during
	failure. Useful for multiruns, saving info about the crash, etc.

	:param cfg: A DictConfig configuration composed by Hydra.
	:return: A tuple with metrics and dict with all instantiated objects.
	"""
	# Pipeline launcher initializes the modules needed for the pipeline to run.
	# It also serves as a context manager, so all resources are properly closed after the pipeline is done.
	with pipeline_launcher(cfg) as pipeline_modules:

	if cfg.get("train"):
	command_line_logger.info("Starting training!")
	pipeline_modules.trainer.fit(
	model=pipeline_modules.model,
	datamodule=pipeline_modules.datamodule,
	ckpt_path=cfg.get("ckpt_path"),
	)
	train_metrics = pipeline_modules.trainer.callback_metrics

	if cfg.get("test"):
	command_line_logger.info("Starting testing!")
	ckpt_path = None
	# Check if a checkpoint callback is available and if it has a best model path.
	# Note that if multiple checkpoint callbacks are used, only the first one will be used
	# to determine the best model path for testing.
	checkpoint_callback = getattr(
	pipeline_modules.trainer, "checkpoint_callback", None
	)
	if checkpoint_callback:
	ckpt_path = getattr(checkpoint_callback, "best_model_path", None)
	if ckpt_path == "":
	ckpt_path = None
	if not ckpt_path:
	command_line_logger.warning(
	"Best checkpoint not found! Using current weights for testing..."
	)
	pipeline_modules.trainer.test(
	model=pipeline_modules.model,
	datamodule=pipeline_modules.datamodule,
	ckpt_path=ckpt_path,
	)
	command_line_logger.info(f"Best ckpt path: {ckpt_path}")

	test_metrics = pipeline_modules.trainer.callback_metrics

	# merge train and test metrics
	metric_dict = {train_metrics, test_metrics}

	command_line_logger.info(f"Metrics: {metric_dict}")


	@hydra.main(version_base="1.3", config_path="../configs", config_name="train.yaml")
	def main(cfg: DictConfig) -> Optional[float]:
	"""Main entry point for training.

	:param cfg: DictConfig configuration composed by Hydra.
	:return: Optional[float] with optimized metric value.
	"""
	# apply extra utilities
	# (e.g. ask for tags if none are provided in cfg, print cfg tree, etc.)
	extras(cfg)
	job_launcher = LocalJobLauncher(cfg=cfg)
	job_launcher.launch(function_to_run=train)


	if __name__ == "__main__":
	main()