Spaces:
Sleeping
Sleeping
File size: 10,647 Bytes
eb9c81a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 | import os
import warnings
from typing import Any, Dict, Optional, Type
from lightning_fabric.utilities.cloud_io import get_filesystem
from pytorch_lightning import LightningModule, Trainer
from pytorch_lightning.cli import (
LightningArgumentParser,
LightningCLI,
LRSchedulerTypeUnion,
ReduceLROnPlateau,
SaveConfigCallback,
)
from pytorch_lightning.loggers import Logger, WandbLogger
from torch.optim import Optimizer
from torch.optim.lr_scheduler import CyclicLR, OneCycleLR
class WandbSaveConfigCallback(SaveConfigCallback):
def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str) -> None:
if self.already_saved:
return
log_dir = trainer.log_dir # this broadcasts the directory
if trainer.logger is not None and trainer.logger.name is not None and trainer.logger.version is not None:
log_dir = os.path.join(log_dir, trainer.logger.name, str(trainer.logger.version))
config_path = os.path.join(log_dir, self.config_filename)
fs = get_filesystem(log_dir)
if not self.overwrite:
# check if the file exists on rank 0
file_exists = fs.isfile(config_path) if trainer.is_global_zero else False
# broadcast whether to fail to all ranks
file_exists = trainer.strategy.broadcast(file_exists)
if file_exists:
raise RuntimeError(
f"{self.__class__.__name__} expected {config_path} to NOT exist. Aborting to avoid overwriting"
" results of a previous run. You can delete the previous config file,"
" set `LightningCLI(save_config_callback=None)` to disable config saving,"
' or set `LightningCLI(save_config_kwargs={"overwrite": True})` to overwrite the config file.'
)
# save the file on rank 0
if trainer.is_global_zero:
# save only on rank zero to avoid race conditions.
# the `log_dir` needs to be created as we rely on the logger to do it usually
# but it hasn't logged anything at this point
fs.makedirs(log_dir, exist_ok=True)
self.parser.save(
self.config, config_path, skip_none=False, overwrite=self.overwrite, multifile=self.multifile
)
self.already_saved = True
# save optimizer and lr scheduler config
# for _logger in trainer.loggers:
# if isinstance(_logger, Logger):
# config = {}
# if "optimizer" in self.config:
# config["optimizer"] = {
# k.replace("init_args.", ""): v for k, v in dict(self.config["optimizer"]).items()
# }
# if "lr_scheduler" in self.config:
# config["lr_scheduler"] = {
# k.replace("init_args.", ""): v for k, v in dict(self.config["lr_scheduler"]).items()
# }
# _logger.log_hyperparams(config)
# save optimizer and lr scheduler config
for _logger in trainer.loggers:
if isinstance(_logger, Logger):
config = {}
optimizer_config = self.config.get("optimizer", None)
if optimizer_config is not None:
if isinstance(optimizer_config, list):
config["optimizer"] = [
{k.replace("init_args.", ""): v for k, v in opt_conf.items()}
for opt_conf in optimizer_config if opt_conf is not None
]
else:
config["optimizer"] = {
k.replace("init_args.", ""): v for k, v in optimizer_config.items()
}
lr_scheduler_config = self.config.get("lr_scheduler", None)
if lr_scheduler_config is not None:
if isinstance(lr_scheduler_config, list):
config["lr_scheduler"] = [
{k.replace("init_args.", ""): v for k, v in sch_conf.items()}
for sch_conf in lr_scheduler_config if sch_conf is not None
]
else:
config["lr_scheduler"] = {
k.replace("init_args.", ""): v for k, v in lr_scheduler_config.items()
}
_logger.log_hyperparams(config)
# broadcast so that all ranks are in sync on future calls to .setup()
self.already_saved = trainer.strategy.broadcast(self.already_saved)
class CustomLightningCLI(LightningCLI):
def __init__(
self,
save_config_callback: Optional[Type[SaveConfigCallback]] = WandbSaveConfigCallback,
parser_kwargs: Optional[Dict[str, Any]] = None,
**kwargs: Any,
) -> None:
new_parser_kwargs = {
sub_command: dict(default_config_files=[os.path.join("configs", "default.yaml")])
for sub_command in ["fit", "validate", "test", "predict"]
}
new_parser_kwargs.update(parser_kwargs or {})
super().__init__(save_config_callback=save_config_callback, parser_kwargs=new_parser_kwargs, **kwargs)
def add_arguments_to_parser(self, parser: LightningArgumentParser) -> None:
parser.add_argument("--ignore_warnings", default=False, type=bool, help="Ignore warnings")
parser.add_argument("--git_commit_before_fit", default=False, type=bool, help="Git commit before training")
parser.add_argument(
"--test_after_fit", default=True, type=bool, help="Run test on the best checkpoint after training"
)
parser.add_argument("--validate_every_n_steps", default=10000, type=int, help="Run validation every N steps (0 to disable)")
parser.add_argument("--test_every_epoch", default=False, type=bool, help="Run test after every epoch")
parser.add_argument(
"--load_weights_path",
type=str,
default=None,
help="If set, load model weights from this checkpoint path (but do NOT load optimizer state)."
)
def before_instantiate_classes(self) -> None:
if self.config[self.subcommand].get("ignore_warnings"):
warnings.filterwarnings("ignore")
def after_instantiate_classes(self) -> None:
"""
This is called after:
- arguments have been parsed,
- your LightningModule (self.model) has been created,
- datamodule (if any) has been created,
- and trainer has been instantiated (but not yet fitted).
We use it to load just the model‐weights if --load_weights_path was provided.
"""
load_path = self.config.get("load_weights_path", None)
if load_path:
ckpt = torch.load(load_path, map_location=lambda storage, loc: storage)
# Only load the model's state_dict; optimizer / lr_scheduler are untouched.
self.model.load_state_dict(ckpt["state_dict"])
def before_fit(self) -> None:
if self.config.fit.get("git_commit_before_fit") and not os.environ.get("DEBUG", False):
logger = self.trainer.logger
if isinstance(logger, WandbLogger):
version = getattr(logger, "version")
name = getattr(logger, "_name")
message = "Commit Message"
if name and version:
message = f"{name}_{version}"
elif name:
message = name
elif version:
message = version
os.system(f'git commit -am "{message}"')
# Setup validation interval if specified
validate_every_n_steps = self.config.fit.get("validate_every_n_steps", 0)
if validate_every_n_steps > 0:
self.trainer.fit_loop.epoch_loop.val_check_batch = validate_every_n_steps
def after_fit(self) -> None:
if self.config.fit.get("test_after_fit") and not os.environ.get("DEBUG", False):
self._run_subcommand("test")
def after_train_epoch(self) -> None:
if self.config.fit.get("test_every_epoch") and not os.environ.get("DEBUG", False):
self._run_subcommand("test")
def before_test(self) -> None:
if self.trainer.checkpoint_callback and self.trainer.checkpoint_callback.best_model_path:
tested_ckpt_path = self.trainer.checkpoint_callback.best_model_path
elif self.config_init[self.config_init["subcommand"]]["ckpt_path"]:
return
else:
tested_ckpt_path = None
self.config_init[self.config_init["subcommand"]]["ckpt_path"] = tested_ckpt_path
def _prepare_subcommand_kwargs(self, subcommand: str) -> Dict[str, Any]:
"""Prepares the keyword arguments to pass to the subcommand to run."""
fn_kwargs = {
k: v
for k, v in self.config_init[self.config_init["subcommand"]].items()
if k in self._subcommand_method_arguments[subcommand]
}
fn_kwargs["model"] = self.model
if self.datamodule is not None:
fn_kwargs["datamodule"] = self.datamodule
return fn_kwargs
@staticmethod
def configure_optimizers(
lightning_module: LightningModule, optimizer: Optimizer, lr_scheduler: Optional[LRSchedulerTypeUnion] = None
) -> Any:
"""Override to customize the :meth:`~pytorch_lightning.core.LightningModule.configure_optimizers` method.
Args:
lightning_module: A reference to the model.
optimizer: The optimizer.
lr_scheduler: The learning rate scheduler (if used).
"""
if lr_scheduler is None:
return optimizer
if isinstance(lr_scheduler, ReduceLROnPlateau):
return {
"optimizer": optimizer,
"lr_scheduler": {"scheduler": lr_scheduler, "monitor": lr_scheduler.monitor},
}
if isinstance(lr_scheduler, (OneCycleLR, CyclicLR)):
# CyclicLR and OneCycleLR are step-based schedulers, where the default interval is "epoch".
return {"optimizer": optimizer, "lr_scheduler": {"scheduler": lr_scheduler, "interval": "step"}}
return [optimizer], [lr_scheduler]
|