Package: hyperstate
HyperStateOpinionated library for managing hyperparameter configs and mutable program state of machine learning training systems.Key Features:(De)serialize nested Python dataclasses asRusty Object NotationOverride any config value from the command lineAutomatic checkpointing and restoration of full program stateCheckpoints are (partially) human readable and can be modified in a text editorPowerful tools for versioning and schema evolution that can detect breaking changes and make it easy to restructure your program while remaining backwards compatible with old checkpointsLarge binary objects in checkpoints can be loaded lazily only when accessedFermented-vegetable freeDSL for hyperparameter schedules(planned) Edit hyperparameters of running experiments on the fly without restartsQuick start guideInstall with pip:pip install hyperstateAll you need to use HyperState is a (nested) dataclass for your hyperparameters:fromdataclassesimportdataclass@dataclassclassOptimizerConfig:lr:float=0.003batch_size:int=512@dataclassclassNetConfig:hidden_size:int=128num_layers:int=2@dataclassclassConfig:optimizer:OptimizerConfignet:NetConfigsteps:int=100Thehyperstate.loadfunction can load values from a config file and/or apply specific overrides from the command line.importargparseimporthyperstateif__name__=="__main__":parser=argparse.ArgumentParser()parser.add_argument("--config",type=str,default=None,help="Path to config file")parser.add_argument("--hps",nargs="+",help="Override hyperparameter value")args=parser.parse_args()config=hyperstate.load(Config,file=args.config,overrides=args.hps)print(config)$pythonmain.py--hpsnet.num_layers=96steps=50Config(optimizer=OptimizerConfig(lr=0.003,batch_size=512),net=NetConfig(hidden_size=128,num_layers=96),steps=50)$catconfig.ron
Config(optimizer:(lr:0.05,batch_size:4096,),)$pythonmain.py--config=config.ron
Config(optimizer=OptimizerConfig(lr=0.05,batch_size=4096),net=NetConfig(hidden_size=128,num_layers=2),steps=100)The full code for this example can be found inexamples/basic-config.Learn more about:ConfigsVersioning and schema evolutionSerializing complex objectsCheckpointing and schedulesExample applicationConfigsHyperState supports a strictly typed subset of Python objects:dataclassescontainers:Dict,List,Tuple,Optionalprimitives:int,float,str,Enumobjects with custom serialization logic:hyperstate.SerializableUsehyperstate.dumpto serialize configs.
The second argument todumpis a path to a file, and can be omitted to return the serialized config as a string instead of saving it to a file:>>>print(hyperstate.dump(Config(lr=0.1,batch_size=256))Config(lr:0.1,batch_size:256,)Usehyperstate.loadto deserialize configs.
Theloadmethod takes the type of the config as the first argugment, and allows you to optionally specify the path to a config file and/or aList[str]of overrides:@dataclassclassOptimizerConfig:lr:floatbatch_size:int@dataclassclassConfig:optimzer:OptimizerConfigsteps:intconfig=hyperstate.load(Config,file="config.ron",overrides=["optimizer.lr=0.1","steps=100"])VersioningVersioning allows you to modify yourConfigclass while still remaining compatible with checkpoints recorded at previous version.
To benefit from versionining, your config must inherithyperstate.Versionedand implement itsversionfunction:@dataclassclassConfig(hyperstate.Versioned):lr:floatbatch_size:int@classmethoddefversion(clz)->int:return0When serializing the config, hyperstate will now record an additionalversionfield with the value of the current version.
Any snapshots that contain configs without a version field are assumed to have a version of0.RewriteRuleNow suppose you modify yourConfigclass, e.g. by renaming thelrfield tolearning_rate.
To still be able to load old configs that are usinglrinstead oflearning_rate, you increase theversionto1and add an entry to the dictionary returned byupgrade_rulesthat tells HyperState to changelrtolearning_ratewhen upgrading configs from version0.fromdataclassesimportdataclassfromtypingimportDict,ListfromhyperstateimportVersionedfromhyperstate.schema.rewrite_ruleimportRenameField,RewriteRule@dataclassclassConfig(Versioned):learning_rate:floatbatch_size:int@classmethoddefversion(clz)->int:return1@classmethoddefupgrade_rules(clz)->Dict[int,List[RewriteRule]]:"""Returns a list of rewrite rules that can be applied to the given versionto make it compatible with the next version."""return{0:[RenameField(old_field=("lr",),new_field=("learning_rate",))],}In the majority of cases, you don't actually have to manually write outRewriteRules.
Instead, they are generated for you automatically by theSchema Evolution CLI.Schema evolution CLIHyperState comes with a command line tool for managing changes to your config schema.
To access the CLI, simply add the following code to the Python file defining your config:# config.pyfromhyperstateimportschema_evolution_cliif__name__=="__main__":schema_evolution_cli(Config)Runpython config.pyto see a list of available commands, described in more detail below.dump-schemaThedump-schemacommand creates a file describing the schema of your config.
This file should commited to version control, and is used to detect changes to the config schema and perform automatic upgrades.check-schemaThecheck-schemacommand compares your config class to a schema file and detects any backwards incompatible changes.
It also emits a suggested list ofRewriteRules that can upgrade old configs to the new schema.
HyperState does not always guess the correctRewriteRules so you still need to check that they are correct.$ python config.py check-schema
WARN  field renamed to learning_rate: lr
WARN  schema changed but version identical
Schema incompatible

Proposed mitigations
- add upgrade rules:
    0: [
        RenameField(old_field=('lr',), new_field=('learning_rate',)),
    ],
- bump version to 1upgrade-schemaTheupgrade-schemacommand functions much the same ascheck-schema, but also updates your schema config files once all backwards-incompatability issues have been address.upgrade-configTheupgrade-configcommand takes a list of paths to config files, and upgrades them to the latest version.Automated TestsTo prevent accidental backwards-incompatible modifications of yourConfigclass, you can use the following code as an automated test that checks your configClassagainst a schema file created withdump-schema:fromhyperstate.schema.schema_changeimportSeverityfromhyperstate.schema.schema_checkerimportSchemaCheckerfromhyperstate.schema.typesimportload_schemafromconfigimportConfigdeftest_schema():old=load_schema("config-schema.ron")checker=SchemaChecker(old,Config)ifchecker.severity()>=Severity.WARN:checker.print_report()assertchecker.severity()==Severity.INFO[unstable feature]SerializableYou can define custom serialization logic for a class by inheriting fromhyperstate.Serializableand implementing theserializeanddeserializemethods.fromdataclassesimportdataclassimporttorchimporttorch.nnasnnimporttorch.nn.functionalasFimporthyperstate@dataclassclassConfig:inputs:intclassLinearRegression(nn.Module,hyperstate.Serializable):def__init__(self,inputs):super(Net,self).__init__()self.fc1=nn.Linear(inputs,1)defforward(self,x):returnself.fc1(x)# `serialize` should return a representation of the object consisting only of# primitives, containers, numpy arrays, and torch tensors.defserialize(self)->Any:returnself.state_dict()# `deserialize` should take a serialized representation of the object and# return an instance of the class. The `ctx` argument allows you to pass# additional information to the deserialization function.@classmethoddefdeserialize(clz,state_dict,ctx):net=clz(ctx["config"].inputs)returnnet.load_state_dict(state_dict)@dataclassclassState:net:LinearRegressionconfig=hyperstate.load("config.ron")state=hyperstate.load("state.ron",ctx={"config":config})Objects that implementSerializableare stored in separate files using a binary encoding.
In the above example, callinghyperstate.dump(state, "checkpoint/state.ron")will result in the following file structure:checkpoint
├── state.net.blob
└── state.ron[unstable feature]LazyIf you inherit fromhyperstate.Lazy, any fields withSerializabletypes will only be loaded/deserialized when accessed. If the.blobfile for a field is missing, HyperState will not raise an error unless the corresponding field is accessed.[unstable feature]blobTo include objects in your state that do not directly implementhyperstate.Serializable, you can seperately implementhyperstate.Serializableand use theblobfunction to mix in theSerializableimplementation:importtorch.optimasoptimimporttorch.nnasnnimporthyperstateclassSerializableOptimizer(hyperstate.Serializable):defserialize(self):returnself.state_dict()@classmethoddefdeserialize(clz,state_dict:Any,config:Config,state:"State")->optim.Optimizer:optimizer=blob(optim.SerializableAdam,mixin=SerializableOptimizer)(state.net.parameters())optimizer.load_state_dict(state_dict)returnoptimizer@dataclassclassState(hyperstate.Lazy):net:nn.Moduleoptimizer:blob(Adam,mixin=SerializableOptimizer)[unstable feature]HyperStateTo unlock the full power of HyperState, you must inherit from theHyperStateclass.
This class combines an immutable config and mutable state, and provides automatic checkpointing, hyperparameter schedules, and the on-the-fly changes to the config and state (not implemented yet).fromdataclassesimportdataclassimporttorchimporttorch.nnasnnimporttorch.nn.functionalasFimporthyperstate@dataclassclassConfig:inputs:intsteps:intclassLinearRegression(nn.Module,hyperstate.Serializable):def__init__(self,inputs):super(Net,self).__init__()self.fc1=nn.Linear(inputs,1)defforward(self,x):returnself.fc1(x)defserialize(self)->Any:returnself.state_dict()@classmethoddefdeserialize(clz,state_dict,ctx):net=clz(ctx["config"].inputs)returnnet.load_state_dict(state_dict)@dataclassclassState:net:LinearRegressionstep:intclassTrainer(HyperState[Config,State]):def__init__(self,# Path to the config fileinitial_config:str,# Optional path to the checkpoint directory, which enables automatic checkpointing.# If any checkpoint files are present, they will be used to initialize the state.checkpoint_dir:Optional[str]=None,# List of manually specified config overrides.config_overrides:Optional[List[str]]=None,):super().__init__(Config,State,initial_config,checkpoint_dir,overrides=config_overrides)definitial_state(self)->State:"""This function is called to initialize the state if no checkpoint files are found."""returnState(net=LinearRegression(self.config.inputs))deftrain(self)->None:forstepinrange(self.state.step,self.config.steps):# training code...self.state.step=step# At the end of each iteration, call `self.step()` to checkpoint the state and apply hyperparameter schedules.self.step()[unstable feature]CheckpointingWhen using theHyperStateobject, the config and state are automatically checkpointed to the configured directory when calling thestepmethod.[unstable feature]SchedulesAnyint/floatfields in the config can also be set to a schedule that will be updated at each step.
For example, the following config defines a schedule that linearly decays the learning rate from 1.0 to 0.1 over 1000 steps:Config(lr:Schedule(key:"state.step",schedule:[(0,1.0),"lin",(1000,0.1),],),batch_size:256,)When you callstep(), all config values that are schedules will be updated.LicenseHyperState is dual-licensed under the MIT license and Apache License (Version 2.0).See LICENSE-MIT and LICENSE-APACHE for more information.