ml.core.registry
Defines the registry for all objects in the project.
The registry is used to register all objects in the project, and to construct them from their configurations. This is done by using the register decorator, which registers the decorated class with an associated config dataclass in the registry. The registry can then be used to construct the object from its configuration.
- ml.core.registry.add_project_dir(path: Path) None
- class ml.core.registry.register_base(name: str, config: type[Config])[source]
Bases:
ABC,Generic[Entry,Config]Defines the base registry type.
- REGISTRY: dict[str, tuple[type[Entry], Union[type[Config], Config]]] = {}
- REGISTRY_LOCATIONS: dict[str, pathlib.Path] = {}
- abstract classmethod config_key() str[source]
Returns the key for the current item from the config.
- classmethod populate_registry(name: str) None[source]
Populates the registry until it has the requested name available.
- Parameters:
name – The name of the registry item to get
- class ml.core.registry.multi_register_base(name: str, config: type[Config])[source]
Bases:
register_base[Entry,Config],Generic[Entry,Config]Defines a registry which produces multiple objects.
- class ml.core.registry.register_model(name: str, config: type[Config])[source]
Bases:
register_base[BaseModel,BaseModelConfig]Defines a registry for holding modules.
- REGISTRY: dict[str, tuple[type['BaseModel'], 'type[BaseModelConfig] | BaseModelConfig']] = {'gan': (<class 'ml.models.gan.GenerativeAdversarialNetworkModel'>, <class 'ml.models.gan.GenerativeAdversarialNetworkModelConfig'>)}
- REGISTRY_LOCATIONS: dict[str, pathlib.Path] = {'gan': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/models/gan.py')}
- class ml.core.registry.register_task(name: str, config: type[Config])[source]
Bases:
register_base[BaseTask,BaseTaskConfig]Defines a registry for holding tasks.
- REGISTRY: dict[str, tuple[type['BaseTask'], 'type[BaseTaskConfig] | BaseTaskConfig']] = {}
- REGISTRY_LOCATIONS: dict[str, pathlib.Path] = {}
- class ml.core.registry.register_trainer(name: str, config: type[Config])[source]
Bases:
register_base[BaseTrainer,BaseTrainerConfig]Defines a registry for holding trainers.
- REGISTRY: dict[str, tuple[type['BaseTrainer'], 'type[BaseTrainerConfig] | BaseTrainerConfig']] = {'gan': (<class 'ml.trainers.gan.GenerativeAdversarialNetworkTrainer'>, <class 'ml.trainers.gan.GenerativeAdversarialNetworkTrainerConfig'>), 'rl': (<class 'ml.trainers.rl.ReinforcementLearningTrainer'>, <class 'ml.trainers.rl.ReinforcementLearningTrainerConfig'>), 'sl': (<class 'ml.trainers.sl.SupervisedLearningTrainer'>, <class 'ml.trainers.sl.SupervisedLearningTrainerConfig'>)}
- REGISTRY_LOCATIONS: dict[str, pathlib.Path] = {'gan': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/trainers/gan.py'), 'rl': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/trainers/rl.py'), 'sl': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/trainers/sl.py')}
- class ml.core.registry.register_optimizer(name: str, config: type[Config])[source]
Bases:
register_base[BaseOptimizer,BaseOptimizerConfig]Defines a registry for holding optimizers.
- REGISTRY: dict[str, tuple[type['BaseOptimizer'], 'type[BaseOptimizerConfig] | BaseOptimizerConfig']] = {'adam': (<class 'ml.optimizers.adam.AdamOptimizer'>, <class 'ml.optimizers.adam.AdamOptimizerConfig'>), 'adan': (<class 'ml.optimizers.adan.AdanOptimizer'>, <class 'ml.optimizers.adan.AdanOptimizerConfig'>), 'gan': (<class 'ml.optimizers.gan.GenerativeAdversarialNetworkOptimizer'>, <class 'ml.optimizers.gan.GenerativeAdversarialNetworkOptimizerConfig'>), 'gpt-3-large': (<class 'ml.optimizers.adam.AdamOptimizer'>, AdamOptimizerConfig(name='???', lr=0.00025, betas=(0.9, 0.95), eps=1e-05, weight_decay=0.1, amsgrad=False, default_decay=True, foreach=None, capturable=False, differentiable=False, fused=None)), 'gpt-3-medium': (<class 'ml.optimizers.adam.AdamOptimizer'>, AdamOptimizerConfig(name='???', lr=0.0003, betas=(0.9, 0.95), eps=1e-05, weight_decay=0.1, amsgrad=False, default_decay=True, foreach=None, capturable=False, differentiable=False, fused=None)), 'gpt-3-small': (<class 'ml.optimizers.adam.AdamOptimizer'>, AdamOptimizerConfig(name='???', lr=0.0006, betas=(0.9, 0.95), eps=1e-05, weight_decay=0.1, amsgrad=False, default_decay=True, foreach=None, capturable=False, differentiable=False, fused=None)), 'lion': (<class 'ml.optimizers.lion.LionOptimizer'>, <class 'ml.optimizers.lion.LionOptimizerConfig'>), 'lion-stable': (<class 'ml.optimizers.lion.LionOptimizer'>, LionOptimizerConfig(name='???', lr=0.0001, betas=(0.95, 0.98), weight_decay=0.01, default_decay=True, use_triton=True)), 'roberta-base': (<class 'ml.optimizers.adam.AdamOptimizer'>, AdamOptimizerConfig(name='???', lr=0.0006, betas=(0.9, 0.98), eps=1e-05, weight_decay=0.01, amsgrad=False, default_decay=True, foreach=None, capturable=False, differentiable=False, fused=None)), 'roberta-large': (<class 'ml.optimizers.adam.AdamOptimizer'>, AdamOptimizerConfig(name='???', lr=0.0004, betas=(0.9, 0.98), eps=1e-05, weight_decay=0.01, amsgrad=False, default_decay=True, foreach=None, capturable=False, differentiable=False, fused=None)), 'sgd': (<class 'ml.optimizers.sgd.SGDOptimizer'>, <class 'ml.optimizers.sgd.SGDOptimizerConfig'>), 'shampoo': (<class 'ml.optimizers.shampoo.ShampooOptimizer'>, <class 'ml.optimizers.shampoo.ShampooOptimizerConfig'>)}
- REGISTRY_LOCATIONS: dict[str, pathlib.Path] = {'adam': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/optimizers/adam.py'), 'adan': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/optimizers/adan.py'), 'gan': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/optimizers/gan.py'), 'gpt-3-large': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/optimizers/adam.py'), 'gpt-3-medium': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/optimizers/adam.py'), 'gpt-3-small': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/optimizers/adam.py'), 'lion': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/optimizers/lion.py'), 'lion-stable': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/optimizers/lion.py'), 'roberta-base': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/optimizers/adam.py'), 'roberta-large': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/optimizers/adam.py'), 'sgd': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/optimizers/sgd.py'), 'shampoo': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/optimizers/shampoo.py')}
- class ml.core.registry.register_lr_scheduler(name: str, config: type[Config])[source]
Bases:
register_base[BaseLRScheduler,BaseLRSchedulerConfig]Defines a registry for holding learning rate schedulers.
- REGISTRY: dict[str, tuple[type['BaseLRScheduler'], 'type[BaseLRSchedulerConfig] | BaseLRSchedulerConfig']] = {'constant': (<class 'ml.lr_schedulers.constant.ConstantLRScheduler'>, <class 'ml.lr_schedulers.constant.ConstantLRSchedulerConfig'>), 'cosine': (<class 'ml.lr_schedulers.cosine.CosineLRScheduler'>, <class 'ml.lr_schedulers.cosine.CosineLRSchedulerConfig'>), 'cosine_decay': (<class 'ml.lr_schedulers.cosine_decay.CosineDecayLRScheduler'>, <class 'ml.lr_schedulers.cosine_decay.CosineDecayLRSchedulerConfig'>), 'gan': (<class 'ml.lr_schedulers.gan.GenerativeAdversarialNetworkLRScheduler'>, <class 'ml.lr_schedulers.gan.GenerativeAdversarialNetworkLRSchedulerConfig'>), 'linear': (<class 'ml.lr_schedulers.linear.LinearLRScheduler'>, <class 'ml.lr_schedulers.linear.LinearLRSchedulerConfig'>), 'linear_no_decay': (<class 'ml.lr_schedulers.linear_no_decay.LinearNoDecayLRScheduler'>, <class 'ml.lr_schedulers.linear_no_decay.LinearNoDecayLRSchedulerConfig'>)}
- REGISTRY_LOCATIONS: dict[str, pathlib.Path] = {'constant': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/lr_schedulers/constant.py'), 'cosine': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/lr_schedulers/cosine.py'), 'cosine_decay': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/lr_schedulers/cosine_decay.py'), 'gan': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/lr_schedulers/gan.py'), 'linear': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/lr_schedulers/linear.py'), 'linear_no_decay': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/lr_schedulers/linear_no_decay.py')}
- class ml.core.registry.register_logger(name: str, config: type[Config])[source]
Bases:
multi_register_base[BaseLogger,BaseLoggerConfig]Defines a registry for holding loggers.
- REGISTRY: dict[str, tuple[type['BaseLogger'], 'type[BaseLoggerConfig] | BaseLoggerConfig']] = {'meter': (<class 'ml.loggers.meter.MeterLogger'>, <class 'ml.loggers.meter.MeterLoggerConfig'>), 'stdout': (<class 'ml.loggers.stdout.StdoutLogger'>, <class 'ml.loggers.stdout.StdoutLoggerConfig'>), 'tensorboard': (<class 'ml.loggers.tensorboard.TensorboardLogger'>, <class 'ml.loggers.tensorboard.TensorboardLoggerConfig'>), 'wandb': (<class 'ml.loggers.wandb.WandbLogger'>, <class 'ml.loggers.wandb.WandbLoggerConfig'>)}
- REGISTRY_LOCATIONS: dict[str, pathlib.Path] = {'meter': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/loggers/meter.py'), 'stdout': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/loggers/stdout.py'), 'tensorboard': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/loggers/tensorboard.py'), 'wandb': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/loggers/wandb.py')}
- class ml.core.registry.register_launcher(name: str, config: type[Config])[source]
Bases:
register_base[BaseLauncher,BaseLauncherConfig]- REGISTRY: dict[str, tuple[type['BaseLauncher'], 'type[BaseLauncherConfig] | BaseLauncherConfig']] = {'mp': (<class 'ml.launchers.mp.MultiProcessLauncher'>, <class 'ml.launchers.mp.MultiProcessLauncherConfig'>), 'slurm': (<class 'ml.launchers.slurm.SlurmLauncher'>, <class 'ml.launchers.slurm.SlurmLauncherConfig'>), 'torchrun': (<class 'ml.launchers.torchrun.TorchRunLauncher'>, <class 'ml.launchers.torchrun.TorchRunLauncherConfig'>)}
- REGISTRY_LOCATIONS: dict[str, pathlib.Path] = {'mp': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/launchers/mp.py'), 'slurm': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/launchers/slurm.py'), 'torchrun': PosixPath('/home/runner/work/ml-starter/ml-starter/ml/launchers/torchrun.py')}
- class ml.core.registry.Objects(raw_config: omegaconf.dictconfig.DictConfig, model: 'BaseModel | None' = None, task: 'BaseTask | None' = None, trainer: 'BaseTrainer | None' = None, optimizer: 'BaseOptimizer | None' = None, lr_scheduler: 'BaseLRScheduler | None' = None, logger: 'list[BaseLogger] | None' = None, launcher: 'BaseLauncher | None' = None)[source]
Bases:
object- raw_config: DictConfig
- trainer: BaseTrainer | None = None
- optimizer: BaseOptimizer | None = None
- lr_scheduler: BaseLRScheduler | None = None
- logger: list[BaseLogger] | None = None
- launcher: BaseLauncher | None = None
- classmethod update_config(config: DictConfig) None[source]
Updates the config in-place.
- Parameters:
config – The config to update
- classmethod resolve_config(config: DictConfig) None[source]
Resolves the config in-place.
- Parameters:
config – The config to resolve.
- classmethod parse_raw_config(config: DictConfig, objs: Objects | None = None, ignore: set[Literal['model', 'task', 'trainer', 'optimizer', 'lr_scheduler', 'logger', 'launcher']] | None = None) Objects[source]
Parses a raw config to the objects it contains.
- Parameters:
config – The raw DictConfig to parse.
objs – Objects which have already been parsed.
ignore – A set of object types to ignore.
- Returns:
The parsed Objects dataclass.