Package: hiraishin
HiraishinA thin PyTorch-Lightning wrapper for building configuration-based DL pipelines with Hydra.DependenciesPyTorch LightningHydraPydanticetc.Installation$pipinstall-UhiraishinBasic workflow1. Model initialization with type annotationsDefine a model class that has training components with type annotations.importtorch.nnasnnimporttorch.optimasoptimfromhiraishin.modelsimportBaseModelclassToyModel(BaseModel):net:nn.Linearcriterion:nn.CrossEntropyLossoptimizer:optim.Adamscheduler:optim.lr_schedulers.ExponentialLRdef__init__(self,config:DictConfig)->None:super().__init__(config)Modules with the following prefixes are instantiated by their own role-specific logic.netcriterionoptimizerschedulerThe same notation can be used to define components other than the learning components listed above (e.g., tokenizers). It is also possible to define built-in type constants that are YAML serializable.classToyModel(BaseModel):net:nn.Linearcriterion:nn.CrossEntropyLossoptimizer:optim.Adamscheduler:optim.lr_schedulers.ExponentialLR# additional components and constantstokenizer:MyTokenizern_classes:intdef__init__(self,config:DictConfig)->None:super().__init__(config)2. Configuration file generationHiraishin provides a CLI command that automatically generates a configuration file based on type annotations.For example, ifToyModelis defined inmodels.py(i.e.,from models import ToyModelcan be executed in the code), then the following command will generate the configuration file automatically.$hiraishingeneratemodel.ToyModel--output_dirconfig/model
Theconfighasbeengenerated!-->config/model/ToyModel.yamlLet's take a look at the generated file._target_:models.ToyModel_recursive_:falseconfig:networks:net:args:_target_:torch.nn.Linearout_features:???in_features:???weights:initializer:nullpath:nulllosses:criterion:args:_target_:torch.nn.CrossEntropyLossweight:1.0optimizers:optimizer:args:_target_:torch.optim.Adamparams:-???scheduler:args:_target_:torch.optim.lr_scheduler.ExponentialLRgamma:???interval:epochfrequency:1strict:truemonitor:nulltokenizer:_target_:MyTokenizern_classes:???First of all, it is compliant with the instantiation byhydra.utils.instantiate.The positional arguments are filled with???that indicates mandatory parameters. They should be overridden by the values you want to set.3. Training routines definitionThe rest of model definition is only defining your training routine along with the style of PyTorch Lightning.classToyModel(BaseModel):...defforward(self,x:torch.Tensor)->torch.Tensor:returnself.net(x)deftraining_step(self,batch,*args,**kwargs)->torch.Tensor:x,target=batchpred=self.forward(x)loss=self.criterion(pred,target)self.log('loss/train',loss)returnloss4. Model InstantiationThe defined model can be instantiated from configuration file. Let's train your models!fromhydra.utilsimportinatantiatefromomegeconfimportOmegaConfdefapp():...config=OmegaConf.load('config/model/toy.yaml')model=inatantiate(config)print(model)# ToyModel(#     (net): Linear(in_features=1, out_features=1, bias=True)#     (criterion): CrossEntropyLoss()# )trainer.fit(model,...)5. Model loadingYou can easily load trained models by using the checkpoints generated by PyTorch Lightning's standard features. Let's test your models!fromhiraishin.utilsimportload_from_checkpointmodel=load_from_checkpoint('path/to/model.ckpt')print(model)# ToyModel(#     (net): Linear(in_features=1, out_features=1, bias=True)#     (criterion): CrossEntropyLoss()# )LicenseHiraishin is licensed under the Apache License, Version 2.0. SeeLICENSEfor the full license text.