Package: micro-config
micro_config.pyan opinionated python dataclass config framework for deep learningI hope this approach to configurations can make your life easier, but if it doesn't, please submit a pull request or an issue and I'll see what I can do. This config system is certaintly under development, so open to any new ideas or suggestions.Installationpip install:pipinstallmicro-configor install from source:Placemicro_config.pyat the root of your project.Repo GuideThe config framework is defined inmicro_config.py.The rest of the repo provides a demo for how one might actually want to usemicro_config.pyin a pytorch deep learning project. Specifically, I implement transformer language model training on wikitext in pytorch. For another example of how to usemicro_config.pyin jax (flax or haiku), see thisjax_v_pytorch repo.To run the demo:navigate to the root directorypip install -r requirements.txtexport PYTHONPATH="$PWD"cd scriptspython train_lm.pyOptionally you can define commandline arguments totrain_lm.pylike:pythontrain_lm.pyepochs=1bsize=16model.transformer_config.hidden_dim=256overview of demo project code:scripts/train_lm.pydefines the training configuration and script execution.base_config.pydefines config schema and defaults for all main config objects:WikiDataConfig,TransformerConfig,LMModelConfig,AdamWConfiggeneral_train_loop.pydefines the config schema and script for training models.src/defines all of the core demo project code.Quick Start / WalkthroughMost demo code in this section is adopted from the demo project provided in the repo.Python dataclasses provide a more natural and flexible config definition interface than.yamlfiles.All config schema should be defined as an instance ofConfigScriptorConfigScriptModeland include a@dataclassdecoratorConfigScripts firstly define a parameter schema and optionally default config values.For example, a simple dataset object configuration:fromdataclassesimportdataclass,adsictfrommicro_configimportConfigScript# data config@dataclassclassWikiDataConfig(ConfigScript):f_path:str='data/wikitext-2-raw/wiki.train.raw'max_len:int=256ConfigScripts load associated objects or functions.To do this, allConfigScripts implementunroll(self, metaconfig: MetaConfig).Themetaconfigparameter is another dataclass which specifies configs for the config framework. Feel free to subclassMetaConfig.For example, loading the dataset from the config:fromdataclassesimportdataclass,adsictfrommicro_configimportConfigScript,MetaConfigfromsrc.dataimportWikitextDatasetimporttorchimportos# data config@dataclassclassWikiDataConfig(ConfigScript):f_path:str='data/wikitext-2-raw/wiki.train.raw'max_len:int=256defunroll(self,metaconfig:MetaConfig):# metaconfig.convert_path converts paths reletive to metaconfig.project_root into absolute pathsreturnWikitextDataset(metaconfig.convert_path(self.f_path),self.max_len)if__name__=="__main__":metaconfig=MetaConfig(project_root=os.path.dirname(__file__),verbose=True)data_config=WikiDataConfig(max_len=512)data=data_config.unroll(metaconfig)Configurations can be defined hierarchically.You can defineConfigScriptsas paremeters of otherConfigScriptsYou can define lists or dictionaries ofConfigScripts as parameters of aConfigScriptby wrapping your list or dict inConfigScriptListorConfigScriptDictrespectively.For example, the LM model config below definesConfigScripts for both a dataset and atransformer_configas parameters:frommicro_configimportMetaConfigfrombase_configsimportConfigScriptModelfromdataclassesimportfieldfromsrc.lmimportLMModelimportos# model config@dataclassclassLMModelConfig(ConfigScriptModel):dataset:WikiDataConfig=field(default_factory=lambda:WikiDataConfig())transformer_config:TransformerConfig=field(default_factory=lambda:TransformerConfig(max_len=256))defunroll(self,metaconfig:MetaConfig):dataset=self.dataset.unroll(metaconfig)transformer_config=self.transformer_config.unroll(metaconfig)returnLMModel(dataset,transformer_config,self.device)if__name__=="__main__":metaconfig=MetaConfig(project_root=os.path.dirname(__file__),verbose=True)model_config=LMModelConfig(checkpoint_path=None,strict_load=True,device='cpu',dataset=WikiDataConfig(f_path='data/wikitext-2-raw/wiki.train.raw',max_len=256),transformer_config=TransformerConfig(max_length=256,heads=12,hidden_dim=768,attn_dim=64,intermediate_dim=3072,num_blocks=12,dropout=0.1))model=model_config.unroll(metaconfig)ConfigScriptModel(not provided withmicro_configout of the box), as used above, is a subclass ofConfigScriptwhich defines some default functionality for loading a pytorch module returned by unroll and placing it on a specified device. You can look insidebase_configs.py(or in thejax_v_pytorch repo) to see examples of how to implement special functionality like this.Configs and scripts are unified: a config is to a script as a script is to a config.unroll(self, metaconfig: MetaConfig)can not only be used to load objects, but also to define script logic.For example, let's define a simple configurable training loop:fromsrc.utilsimportcombine_logsfrommicro_configimportConfigScript,MetaConfigfrombase_configsimportConfigScriptModel@dataclassclassTrainLoop(ConfigScript):train_dataset:ConfigScripteval_dataset:ConfigScriptmodel:ConfigScriptModeloptim:ConfigScriptepochs:int=10bsize:int=32defunroll(self,metaconfig:MetaConfig):print('using config:',asdict(self))device=metaconfig.devicetrain_dataset=self.train_dataset.unroll(metaconfig)eval_dataset=self.eval_dataset.unroll(metaconfig)model=self.model.unroll(metaconfig)model.train()train_dataloader=DataLoader(train_dataset,batch_size=self.bsize)eval_dataloader=DataLoader(eval_dataset,batch_size=self.bsize)optim=self.optim.unroll(metaconfig)(model)forepochinrange(epochs):forxintqdm(train_dataloader):loss,logs=model.get_loss(x.to(device))optim.zero_grad()loss.backward()optim.step()model.eval()val_x=next(iter(eval_dataloader))_,val_logs=model.get_loss(val_x.to(device))out_log=print({'train':combine_logs([logs]),'val':combine_logs([val_logs]),'step':(step+1)})model.train()returnmodelObjects returned byunroll(self, metaconfig: MetaConfig)respect the reference structure of the config hierarchy.If the same config object is referenced multiple times in a config hierarchy, the object'sunroll(self, metaconfig: MetaConfig)method will only be called once and its output cached, subsequent calls will return the cached output. If you don't want this caching behavior, you can subclassConfigScriptNoCacheinstead.For example,train_datasetis referenced twice intrain_config_script:importtorchimportostrain_dataset=WikiDataConfig(f_path='data/wikitext-2-raw/wiki.train.raw',max_len=256)eval_dataset=WikiDataConfig(f_path='data/wikitext-2-raw/wiki.valid.raw',max_len=256)model=LMModelConfig(checkpoint_path=None,strict_load=True,device='cpu',dataset=train_dataset,transformer_config=TransformerConfig(max_length=256,heads=12,hidden_dim=768,attn_dim=64,intermediate_dim=3072,num_blocks=12,dropout=0.1))train_config_script=TrainLoop(train_dataset=train_dataset,eval_dataset=eval_dataset,model=model,optim=AdamWConfig(lr=1e-4,weight_decay=0.01),epochs=10,bsize=16,)if__name__=="__main__":metaconfig=MetaConfig(project_root=os.path.dirname(__file__),verbose=True)# run the scripttrain_config_script.unroll(metaconfig)The dataset object configured bytrain_datasetwill only be loaded once in the above hiararchy, even though bothLMModelConfigandTrainLooptake it in as input.A method for parsing commandline args is provided.parse_args(config)parses the command line arguments into a dictionarydeep_replace(config, **kwargs)implements a nested version of the standarddataclasses.replacefunctionfrommicro_configimportparse_args,deep_replace,MetaConfigimportosif__name__=="__main__":metaconfig=MetaConfig(project_root=os.path.dirname(__file__),verbose=True)train_config_script=deep_replace(train_config_script,**parse_args())# run the scripttrain_config_script.unroll(metaconfig)To edit any arguments in the hierarchy through the commandline, call the script like so:pythontrain_lm.pyepochs=1bsize=16model.transformer_config.hidden_dim=256