riccardo / Save and load hyperparameters

from datetime import datetime from pathlib import Path import yaml # Define the model class MyModel: def __init__( self, n_layers: int, n_heads: int, dropout: float, outdir: str | Path, ): self.n_layers = n_layers self.n_heads = n_heads self.dropout = dropout self.outdir = outdir @classmethod def from_config(cls, config: Path | str): with open(config, "r") as f: hparams = yaml.safe_load(f) return cls(**hparams) def load_checkpoint(self, checkpoint: Path | str): print(f"Loading checkpoint from {checkpoint}...") def train(self): print("Training...") print("Saving checkpoint...") with open(Path(self.outdir) / "checkpoint.pt", "w") as f: f.write("Hello, world!") def evaluate(self): print("Evaluating...") # Create the output directory now = datetime.now().strftime("%Y%m%d_%H%M%S") outdir = Path(f"{now}_runs") outdir.mkdir(parents=True, exist_ok=True) # Instantiate the model hparams = { "n_layers": 3, "n_heads": 8, "dropout": 0.1, "outdir": str(outdir), } model = MyModel(**hparams) # Train the model model.train() # Save the hyperparameters in the same directory with open(outdir / "hparams.yaml", "w") as f: yaml.dump(hparams, f) # Evaluate the model new_model = MyModel.from_config(outdir / "hparams.yaml") new_model.load_checkpoint(outdir / "checkpoint.pt") new_model.evaluate()