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 DictCo
(cfg: DictConfig)
| 48 | |
| 49 | @task_wrapper |
| 50 | def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: |
| 51 | """Trains the model. Can additionally evaluate on a testset, using best weights obtained during |
| 52 | training. |
| 53 | |
| 54 | This method is wrapped in optional @task_wrapper decorator, that controls the behavior during |
| 55 | failure. Useful for multiruns, saving info about the crash, etc. |
| 56 | |
| 57 | :param cfg: A DictConfig configuration composed by Hydra. |
| 58 | :return: A tuple with metrics and dict with all instantiated objects. |
| 59 | """ |
| 60 | # set seed for random number generators in pytorch, numpy and python.random |
| 61 | if cfg.get("seed"): |
| 62 | L.seed_everything(cfg.seed, workers=True) |
| 63 | |
| 64 | log.info( |
| 65 | f"Setting `float32_matmul_precision` to {cfg.model.cfg.task.float32_matmul_precision}." |
| 66 | ) |
| 67 | torch.set_float32_matmul_precision(precision=cfg.model.cfg.task.float32_matmul_precision) |
| 68 | |
| 69 | log.info(f"Instantiating datamodule <{cfg.data._target_}>") |
| 70 | datamodule: LightningDataModule = hydra.utils.instantiate(cfg.data, stage="fit") |
| 71 | |
| 72 | log.info(f"Instantiating model <{cfg.model._target_}>") |
| 73 | model: LightningModule = hydra.utils.instantiate(cfg.model) |
| 74 | |
| 75 | log.info("Instantiating callbacks...") |
| 76 | callbacks: List[Callback] = instantiate_callbacks(cfg.get("callbacks")) |
| 77 | |
| 78 | log.info("Instantiating loggers...") |
| 79 | logger: List[Logger] = instantiate_loggers(cfg.get("logger")) |
| 80 | |
| 81 | plugins = None |
| 82 | if "_target_" in cfg.environment: |
| 83 | log.info(f"Instantiating environment <{cfg.environment._target_}>") |
| 84 | plugins: ClusterEnvironment = hydra.utils.instantiate(cfg.environment) |
| 85 | |
| 86 | strategy = getattr(cfg.trainer, "strategy", None) |
| 87 | if "_target_" in cfg.strategy: |
| 88 | log.info(f"Instantiating strategy <{cfg.strategy._target_}>") |
| 89 | strategy: Strategy = hydra.utils.instantiate(cfg.strategy) |
| 90 | if ( |
| 91 | "mixed_precision" in strategy.__dict__ |
| 92 | and getattr(strategy, "mixed_precision", None) is not None |
| 93 | ): |
| 94 | strategy.mixed_precision.param_dtype = ( |
| 95 | resolve_omegaconf_variable(cfg.strategy.mixed_precision.param_dtype) |
| 96 | if getattr(cfg.strategy.mixed_precision, "param_dtype", None) is not None |
| 97 | else None |
| 98 | ) |
| 99 | strategy.mixed_precision.reduce_dtype = ( |
| 100 | resolve_omegaconf_variable(cfg.strategy.mixed_precision.reduce_dtype) |
| 101 | if getattr(cfg.strategy.mixed_precision, "reduce_dtype", None) is not None |
| 102 | else None |
| 103 | ) |
| 104 | strategy.mixed_precision.buffer_dtype = ( |
| 105 | resolve_omegaconf_variable(cfg.strategy.mixed_precision.buffer_dtype) |
| 106 | if getattr(cfg.strategy.mixed_precision, "buffer_dtype", None) is not None |
| 107 | else None |
no test coverage detected