diff options
Diffstat (limited to 'training')
| -rw-r--r-- | training/experiments/image_transformer.yaml | 13 | ||||
| -rw-r--r-- | training/run_experiment.py | 50 |
2 files changed, 46 insertions, 17 deletions
diff --git a/training/experiments/image_transformer.yaml b/training/experiments/image_transformer.yaml index 9e8f9fc..bedcbb5 100644 --- a/training/experiments/image_transformer.yaml +++ b/training/experiments/image_transformer.yaml @@ -1,9 +1,12 @@ +seed: 4711 + network: + desc: null type: ImageTransformer args: encoder: - type: None - args: None + type: null + args: null num_decoder_layers: 4 hidden_dim: 256 num_heads: 4 @@ -12,6 +15,7 @@ network: transformer_activation: glu model: + desc: null type: LitTransformerModel args: optimizer: @@ -31,11 +35,11 @@ model: weight: None ignore_index: -100 reduction: mean - monitor: val_loss mapping: sentence_piece data: + desc: null type: IAMExtendedParagraphs args: batch_size: 16 @@ -55,6 +59,7 @@ callbacks: patience: 10 trainer: + desc: null args: stochastic_weight_avg: true auto_scale_batch_size: binsearch @@ -62,6 +67,6 @@ trainer: fast_dev_run: false gpus: 1 precision: 16 - max_epocs: 512 + max_epochs: 512 terminate_on_nan: true weights_summary: true diff --git a/training/run_experiment.py b/training/run_experiment.py index 289866e..1fb4bc2 100644 --- a/training/run_experiment.py +++ b/training/run_experiment.py @@ -6,7 +6,7 @@ from typing import Dict, List, NamedTuple, Optional, Union, Type import click from loguru import logger -from omegaconf import OmegaConf +from omegaconf import DictConfig, OmegaConf import pytorch_lightning as pl import torch from torch import nn @@ -42,6 +42,14 @@ def _configure_logging(log_dir: Optional[Path], verbose: int = 0) -> None: ) +def _load_config(file_path: Path) -> DictConfig: + """Return experiment config.""" + logger.info(f"Loading config from: {file_path}") + if not file_path.exists(): + raise FileNotFoundError(f"Experiment config not found at: {file_path}") + return OmegaConf.load(file_path) + + def _import_class(module_and_class_name: str) -> type: """Import class from module.""" module_name, class_name = module_and_class_name.rsplit(".", 1) @@ -50,26 +58,25 @@ def _import_class(module_and_class_name: str) -> type: def _configure_callbacks( - args: List[Union[OmegaConf, NamedTuple]] + callbacks: List[DictConfig], ) -> List[Type[pl.callbacks.Callback]]: """Configures lightning callbacks.""" pl_callbacks = [ - getattr(pl.callbacks, callback.type)(**callback.args) for callback in args + getattr(pl.callbacks, callback.type)(**callback.args) for callback in callbacks ] return pl_callbacks def _configure_logger( - network: Type[nn.Module], args: Dict, use_wandb: bool + network: Type[nn.Module], args: Dict, use_wandb: bool ) -> pl.loggers.WandbLogger: """Configures lightning logger.""" if use_wandb: pl_logger = pl.loggers.WandbLogger() pl_logger.watch(network) pl_logger.log_hyperparams(vars(args)) - else: - pl_logger = pl.logger.TensorBoardLogger("training/logs") - return pl_logger + return pl_logger + return pl.logger.TensorBoardLogger("training/logs") def _save_best_weights( @@ -89,7 +96,9 @@ def _save_best_weights( wandb.save(best_model_path) -def _load_lit_model(lit_model_class: type, network: Type[nn.Module], config: OmegaConf) -> Type[pl.LightningModule]: +def _load_lit_model( + lit_model_class: type, network: Type[nn.Module], config: DictConfig +) -> Type[pl.LightningModule]: """Load lightning model.""" if config.load_checkpoint is not None: logger.info( @@ -101,8 +110,17 @@ def _load_lit_model(lit_model_class: type, network: Type[nn.Module], config: Ome return lit_model_class(network=network, **config.model.args) -def run(path: str, train: bool, test: bool, tune: bool, use_wandb: bool) -> None: +def run( + filename: str, + train: bool, + test: bool, + tune: bool, + use_wandb: bool, + verbose: int = 0, +) -> None: """Runs experiment.""" + + _configure_logging(None, verbose=verbose) logger.info("Starting experiment...") # Seed everything in the experiment. @@ -110,8 +128,8 @@ def run(path: str, train: bool, test: bool, tune: bool, use_wandb: bool) -> None pl.utilities.seed.seed_everything(SEED) # Load config. - logger.info(f"Loading config from: {path}") - config = OmegaConf.load(path) + file_path = EXPERIMENTS_DIRNAME / filename + config = _load_config(file_path) # Load classes. data_module_class = _import_class(f"text_recognizer.data.{config.data.type}") @@ -169,8 +187,14 @@ def cli( verbose: int, ) -> None: """Run experiment.""" - _configure_logging(None, verbose=verbose) - run(path=experiment_config, train=train, test=test, tune=tune, use_wandb=use_wandb) + run( + filename=experiment_config, + train=train, + test=test, + tune=tune, + use_wandb=use_wandb, + verbose=verbose, + ) if __name__ == "__main__": |