From 732708e0c2c7ab6aa747463ac9e37a540eee7f48 Mon Sep 17 00:00:00 2001 From: Thomas M Kehrenberg Date: Thu, 16 May 2024 20:52:31 +0200 Subject: [PATCH] Add an explanation of experiments and other things --- README.md | 95 +++++++++++++++++++++++++++++ conf/config.yaml | 1 + conf/dm/celeba/male_blond.yaml | 1 + conf/dm/celeba/male_smiling.yaml | 1 + conf/experiment/cmnist_cnn.yaml | 12 ++++ conf/experiment/good_model.yaml | 15 +++++ conf/model/single_linear_layer.yaml | 1 + main.py | 6 +- pyproject.toml | 1 + src/models.py | 2 +- src/optimisation.py | 24 ++++++++ src/run.py | 12 ++-- 12 files changed, 163 insertions(+), 8 deletions(-) create mode 100644 conf/experiment/cmnist_cnn.yaml create mode 100644 conf/experiment/good_model.yaml create mode 100644 src/optimisation.py diff --git a/README.md b/README.md index bfb4b2c..fc7b5fa 100644 --- a/README.md +++ b/README.md @@ -46,12 +46,17 @@ python main.py --help ``` It should look like this: + +
+ Help output (quite long) + ``` == Configuration groups == Compose your configuration from those groups (group=option) dm: celeba, cmnist dm/celeba: male_blond, male_smiling +experiment: cmnist_cnn, good_model model: cnn, fcn, single_linear_layer @@ -78,6 +83,11 @@ model: dropout_prob: 0.0 final_bias: true input_norm: false +opt: + lr: 5.0e-05 + weight_decay: 0.0 + optimizer_cls: torch.optim.AdamW + optimizer_kwargs: null wandb: name: null dir: ./local_logging @@ -94,6 +104,8 @@ seed: 42 gpu: 0 ``` +
+ We can see that `seed` and `gpu` are "top-level" config values, but there are also a lot of values in subconfigs. For example, we can change the `num_hidden` value in the `model` config: ```bash python main.py model.num_hidden=2 @@ -149,6 +161,12 @@ You can also set python main.py dm=celeba/male_smiling ``` +### W&B config +To enable W&B logging, you can set the `wandb.mode` to "online" or "offline" (default is "disabled"): +```bash +python main.py wandb.mode=online +``` + ## Multiruns Often, we want to run the same experiment with only slightly different configurations. This can be done with the `--multirun` flag. For example, to run the code with seeds 42 and 43, we can do: ```bash @@ -186,6 +204,82 @@ By default, Hydra simply runs the code with the different config values sequenti python main.py --multirun seed=42,43 hydra/launcher=slurm/kyiv ``` +Hydra doesn't limit you to iterating over just one parameter. You can also iterate over multiple parameters. For example, to run the code with seeds 42 and 43 and `model.num_hidden` set to 1, 2 and 3, you can do: +```bash +python main.py --multirun seed=42,43 model.hidden_dim=10 model.num_hidden=1,2,3 gpu=0 +``` +This will start 6 jobs: 2 seeds x 3 `model.num_hidden` values. + +## Experiment configs +Hydra commands can get very long when we specify many config values. To make this easier, we can define "experiment configs" in the `conf/experiment/` directory. + +For example, let's say that after a lot of experimentation, we found that the following config values work well together: +```bash +python main.py model=fcn model.num_hidden=2 model.hidden_dim=10 model.norm=LN model.input_norm=true model.activation=SELU opt.lr=0.001 opt.weight_decay=0.001 +``` + +We can then create a new experiment config file at `conf/experiment/good_model.yaml` with the following content: +```yaml +# @package _global_ +--- +defaults: + - override /model: fcn + +model: + num_hidden: 2 + hidden_dim: 10 + norm: LN + input_norm: true + activation: SELU + +opt: + lr: 0.001 + weight_decay: 0.001 +``` + +Note that the comment `# @package _global_` is required. (The reason is that, by default, if you have a config file in the `conf/experiment/` directory, Hydra will want to associate this with the `experiment` entry in the main configuration – which doesn't exist! So, `@package _global_` tells Hydra to put the content of the file at the *top level* of the main config.) + +And then we can run the code with these config values by running: +```bash +python main.py +experiment=good_model +``` + +You can still override values: +```bash +python main.py +experiment=good_model model.hidden_dim=20 +``` + +## How to make your code easily configurable with Hydra +What we found is the best method to structure your code to make it easy to configure with Hydra is the "builder" or "factory" pattern. + +This means you have a dataclass that contains all the configuration values for a particular component of your code (e.g. the model, the data, the optimiser, etc.). And then this class has a `build()` or `init()` method that takes additional arguments which are only available at runtime (e.g. the input size of the model, the number of classes in the dataset, etc.), and then instantiates the component. + +For example, in this code base, we have the `ModelFactory` class in `src/model.py`: +```python +@dataclass(eq=False) +class ModelFactory(ABC): + """Interface for model factories.""" + + @abstractmethod + def build(self, in_dim: int, *, out_dim: int) -> nn.Module: + raise NotImplementedError() +``` +And when you add a model to the code base, you subclass `ModelFactory` and implement the `build()` method: +```python +@dataclass(eq=False, kw_only=True) +class SimpleCNNFactory(ModelFactory): + """Factory for a very simple CNN.""" + + kernel_size: int = 5 + pool_stride: int = 2 + activation: Activation = Activation.RELU + + @override + def build(self, in_dim: int, *, out_dim: int) -> nn.Sequential: + ... +``` + + ## Structure of the code - `main.py`: The main entry point of the code. It sets up Hydra and then calls the main `run()` function. @@ -193,6 +287,7 @@ python main.py --multirun seed=42,43 hydra/launcher=slurm/kyiv - `run.py`: Contains the main `Config` class that is used to define valid config values. It also contains the `run()` function that is called by `main.py`. - `data.py`: Contains the `DataModule` class that is used to load the data. - `model.py`: Contains the `ModelFactory` class that is used to create the model. + - `optimisation.py`: Contains the `OptimisationCfg` class that is used to build the optimiser that trains the model. - `logging.py`: Contains the `WandbCfg` class that is used to set up Weights & Biases logging. - `conf/` - `config.yaml`: The main config file for the project. It sets the default values for `dm` and `model`. diff --git a/conf/config.yaml b/conf/config.yaml index c754969..6ab2beb 100644 --- a/conf/config.yaml +++ b/conf/config.yaml @@ -1,3 +1,4 @@ +--- defaults: # The first entry here always has to be the name we have passed to `register_hydra_config`. - config_schema diff --git a/conf/dm/celeba/male_blond.yaml b/conf/dm/celeba/male_blond.yaml index 42d1665..2904b6c 100644 --- a/conf/dm/celeba/male_blond.yaml +++ b/conf/dm/celeba/male_blond.yaml @@ -1,3 +1,4 @@ +--- defaults: # Specify here the name of the entry in the `dm` config group you want to extend. - celeba diff --git a/conf/dm/celeba/male_smiling.yaml b/conf/dm/celeba/male_smiling.yaml index 113b943..e057a40 100644 --- a/conf/dm/celeba/male_smiling.yaml +++ b/conf/dm/celeba/male_smiling.yaml @@ -1,3 +1,4 @@ +--- defaults: # Specify here the name of the entry in the `dm` config group you want to extend. - celeba diff --git a/conf/experiment/cmnist_cnn.yaml b/conf/experiment/cmnist_cnn.yaml new file mode 100644 index 0000000..2d40950 --- /dev/null +++ b/conf/experiment/cmnist_cnn.yaml @@ -0,0 +1,12 @@ +# @package _global_ +--- +defaults: + - override /dm: cmnist + - override /model: cnn + +dm: + num_colors: 10 + val_prop: 0.2 + +model: + kernel_size: 5 diff --git a/conf/experiment/good_model.yaml b/conf/experiment/good_model.yaml new file mode 100644 index 0000000..7511755 --- /dev/null +++ b/conf/experiment/good_model.yaml @@ -0,0 +1,15 @@ +# @package _global_ +--- +defaults: + - override /model: fcn + +model: + num_hidden: 2 + hidden_dim: 10 + norm: LN + input_norm: true + activation: SELU + +opt: + lr: 0.001 + weight_decay: 0.001 diff --git a/conf/model/single_linear_layer.yaml b/conf/model/single_linear_layer.yaml index 409b4de..b8b54d1 100644 --- a/conf/model/single_linear_layer.yaml +++ b/conf/model/single_linear_layer.yaml @@ -1,3 +1,4 @@ +--- defaults: - fcn diff --git a/main.py b/main.py index a1c6f7a..7cf37a5 100644 --- a/main.py +++ b/main.py @@ -35,7 +35,9 @@ def main(hydra_config: omegaconf.DictConfig) -> None: if __name__ == "__main__": # Before calling the main function, we need to register the main `Config` class and - # the configuration groups. - # Without this, hydra doesn't know which keys and values are valid in the configuration. + # the configuration groups. Without this, hydra doesn't know which keys and values are valid in + # the configuration. + # Whatever you set here as `schema_name` will need to be incluced as the first entry in the + # `defaults` list in the main config yaml file (`conf/config.yaml`). register_hydra_config(Config, CONFIG_GROUPS, schema_name="config_schema") main() diff --git a/pyproject.toml b/pyproject.toml index 32ca81d..e6ffa2a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ pythonVersion = "3.10" typeCheckingMode = "strict" venvPath = "." venv = ".venv" +reportUnknownMemberType = "none" [tool.ruff] line-length = 100 diff --git a/src/models.py b/src/models.py index 3df2518..0cc9102 100644 --- a/src/models.py +++ b/src/models.py @@ -95,7 +95,7 @@ class SimpleCNNFactory(ModelFactory): activation: Activation = Activation.RELU @override - def build(self, in_dim: int, *, out_dim: int) -> nn.Module: + def build(self, in_dim: int, *, out_dim: int) -> nn.Sequential: return nn.Sequential( nn.Conv2d(in_channels=in_dim, out_channels=6, kernel_size=self.kernel_size), self.activation.init(), diff --git a/src/optimisation.py b/src/optimisation.py new file mode 100644 index 0000000..085242d --- /dev/null +++ b/src/optimisation.py @@ -0,0 +1,24 @@ +"""Optimisation and training code.""" + +from collections.abc import Iterator +from dataclasses import dataclass +from typing import Any + +from torch.nn import Parameter +from torch.optim import Optimizer + +__all__ = ["OptimisationCfg"] + + +@dataclass +class OptimisationCfg: + """Config class for the optimisation.""" + + lr: float = 5.0e-5 + weight_decay: float = 0.0 + optimizer_cls: str = "torch.optim.AdamW" + optimizer_kwargs: dict[str, Any] | None = None + + def build(self, params: Iterator[tuple[str, Parameter]]) -> Optimizer: + # Instantiate the optimizer for the given parameters. + raise NotImplementedError() diff --git a/src/run.py b/src/run.py index 7c972f7..251cf32 100644 --- a/src/run.py +++ b/src/run.py @@ -8,6 +8,7 @@ from src.datasets import CelebADataModule, ColoredMNISTDataModule, DataModule from src.logging import WandbCfg from src.models import FcnFactory, ModelFactory, SimpleCNNFactory +from src.optimisation import OptimisationCfg __all__ = ["Config", "CONFIG_GROUPS"] @@ -30,8 +31,9 @@ class Config: dm: DataModule model: ModelFactory - # This is a normal subconfig, for which we can specify a default, - # but note that in dataclasses, the default cannot be mutable, so we use `default_factory`. + # These are normal subconfigs, for which we can specify defaults, + # but note that in dataclasses, the default may not be mutable, so we use `default_factory`. + opt: OptimisationCfg = field(default_factory=OptimisationCfg) wandb: WandbCfg = field(default_factory=WandbCfg) # These are normal fields, for which we can specify defaults. @@ -45,9 +47,9 @@ def run(self, config_for_logging: dict[str, Any]) -> None: torch.manual_seed(self.seed) # Initialize the logger. - run = self.wandb.init(config_for_logging, reinit=True) - if run is not None: - run.log({"accuracy": 0.5}) + wandb_run = self.wandb.init(config_for_logging, reinit=True) + if wandb_run is not None: + wandb_run.log({"accuracy": 0.5}) # Prepare the data module. self.dm.prepare(seed=self.seed)