Package: lightning-cv
Cross validation withLightning FabricInstallationpipinstalllightning-cvPurposePyTorch Lightningprovides many boilerplate-free abstractions to building scalable deep learning models. However, there is little first-class support for cross validation, which is very common when training models with many hyperparameters and heterogeneous data.Fortunately, the creators ofPyTorch Lightninghave developed and released a more flexible package that underliesPyTorch LightningcalledLightning Fabric.Fabricautomatically handles lower technical details like target training device, training precision, and distributed training. BUTFabricleaves the flexibility of setting up training/inference loops, which is beneficial for developing a general purpose cross validation scheme.Overall workflowThe cross validation logic can be described with this pseudocode:importtorchn_folds=len(folds)models=[(MyModel(),MyOptimizer())for_inrange(n_folds)]# tensor that is n_folds large and keeps track of per-fold lossper_fold_loss=torch.zeros((n_folds,))forepochinrange(epochs):forfold,(fold_train_loader,fold_val_loader)inenumerate(folds):model,optimizer=models[fold]# train loop -> step levelforbatchinfold_train_loader:# forward passloss=model(batch)# backward pass and optimizer updateloss.backward()optimizer.step()optimizer.zero_grad()# val loop -> step levelforbatchinfold_val_loader:loss=model(batch)# keep running total at the fold levelper_fold_loss[fold]+=loss# take avg of val loss for current foldper_fold_loss[fold]/=len(fold_val_loader)# at this point, all folds have been trained AND validated once# therefore, we can compute a current val loss as the avg of the# per-fold losscurrent_val_loss=per_fold_loss.mean()# this is useful since we can monitor cross val performance in real-time# here, the report function can do several things:#   1. log performance#   2. model checkpointing#   3. early stopping#   4. make use of a schedulerreport(current_val_loss)# reset per_fold_lossper_fold_loss.fill_(0.0)DetailsFold SynchronizationNotably, this cross validation schemesynchronizes all model folds at the epoch-level. Unfortunately, this requiresk-fold models to be held in memory simultaneously, which might be prohibitive for extremely large models. While looping over the training folds in the outermost loop could be optionally added and reduce the memory footprint, the current cross validation scheme offers greater flexibility.The simplest benefit is that an individual cross validation run will be faster by a factor of ~ksince each model fold is being trained within the epoch-level loop. However, a greater benefit of this approach is that this allows for early stopping of unpromising trials since the current validation metrics can be updated in real-time at the end of each epoch.Epoch level vs step level monitoringSyncing folds at the epoch level means that techniques thatmeaningfullymonitor current model performancemustonly operate at theepochlevel. For example, schedulers can adjust the learning rate based on predefined strategies. The scheduler adjustment frequency, in normal workflows, can happen at the epoch or batch-step level. However, for this cross validation scheme, to ensure more comparable analysis across model folds, schedulers are only allowed to operate at the epoch level, so the learning rate is the same across all model folds.A more explicit way to state this is that each training fold is not necessarily the same size, meaning that each training foldcouldtake a different number of training steps. This becomes problematic when trying to monitor performance at the step level where the model folds are not synced.The exception to this rule is when current performance is read without causing any meaningful updates. This occurs during metric logging and when a progress bar is used.UsageDataThe biggest change that needs to be made to existing workflows is to use alightning.LightningDataModulesubclass that has the methodtrain_val_dataloaders. This method yields the train/val splits of the original dataset.Convenience classes have been provided by this library:fromtypingimportAny,Callable,OptionalimportlightningasLfromtorch.utils.dataimportDatasetfromlightning_cv._typingimportCrossValidatorTypeclass_CrossValidationDataModule(L.LightningDataModule):def__init__(self,dataset:Dataset,batch_size:int,# pass a class type that implements the sklearn split API# ie calling the instance's .split() method yields train/val indicescross_validator:type[CrossValidatorType],# any init kwargs required to instantiate the cross_validatorcross_validator_kwargs:dict[str,Any],# optional torch.utils.data.DataLoader collate_fncollate_fn:Optional[Callable]=None,):super().__init__()# store other input args as attrsself.dataset=datasetself.batch_size=batch_sizeself._cross_validator=cross_validatorself._cross_validator_kwargs=cross_validator_kwargsself.collate_fn=collate_fn# this method can be overridden for a general-use datamoduledefsetup(self,stage):# cross validation only needs the fit stageifstage=="fit":self.data_manager=self._cross_validator(**self._cross_validator_kwargs)deftrain_val_dataloaders(self):fortrain_idx,val_idxinself.data_manager.split():# convert train/val data indices into a DataLoader# actual impl depends on original datasetyieldtrain_loader,val_loaderModelThe model requirements are:RequirementMethodProvided bylightning.LightningModuletraining_stepYYvalidation_stepYYconfigure_optimizersYYlr_scheduler_stepYYtrainYYevalYYfabricNY (technically, still need to set attr)estimated_stepsNNExtend existing modelsMost of these method or attribute requirements are provided by creating models that are subclasses of thelightning.LightningModuleclass. A mixin class is provided that can handle thefabricandestimated_stepsattributes to extend an existingLightningModule.TheCrossValModuleMixinprovides theestimated_stepsproperty and ensures that an instance oflightning.Fabricis passed.TODO: mention pydantic model configimport lightning as L+ from lightning_cv import CrossValModuleMixin- class MyModel(L.LightningModule)+ class MyModel(L.LightningModule, CrossValModuleMixin):def __init__(self, ...):-       super().__init__(self)+       L.LightningModule.__init__(self, ...)+       CrossValModuleMixin.__init__(self, ...)This way your model can be used with either the fully abstractedlightning.Traineror thelightning_cv.CrossValidationTrainer.Create new modelsAlternatively, if you are developing a new model from scratch, you can just subclassCrossValModule, which already subclasseslightning.LightningModuleand theCrossValModuleMixin:fromlightning_cvimportCrossValModule:classMyModel(CrossValModule):def__init__(self,...):# handles L.LightningModule init and# provides other api requirementssuper().__init__(...)# custom init logic hereTrainingConfigTheCrossValidationTrainercanonlybe instantiated with apydanticconfig model.pydanticmodels provide automatic type validation and safety guarantees that are useful for general purpose computing.The config model looks like this. Most of these arguments get passed directly tolightning.Fabric, but several are general model training parameters that need to be tracked.frompydanticimportBaseModelclassCrossValidationTrainerConfig(BaseModel):accelerator:Accelerators|Accelerator="auto"strategy:Strategies|Strategy="auto"devices:list[int]|str|int="auto"precision:Precision|int="32"plugins:Optional[str|Any]=Nonecallbacks:Optional[list[Callback]|Callback]=Noneloggers:Optional[Logger|list[Logger]]=Nonemax_epochs:int=1000grad_accum_steps:int=1limit_train_batches:Number=1.0limit_val_batches:Number=1.0validation_frequency:int=1use_distributed_sampler:bool=Truecheckpoint_dir:Path=Path.cwd().joinpath("checkpoints")checkpoint_frequency:int=1monitor:str="val_loss"You can instantiate the config model like this:fromlightning_cvimportCrossValidationTrainerConfig# use defaultstrainer_config=CrossValidationTrainerConfig()# change defaultstrainer_config=CrossValidationTrainerConfig(max_epochs=2)Trainer classThe trainer class (CrossValidationTrainer) only accepts two init arguments:
a modeltype, and the config instance.Furthemore, the main method on the trainer is thetrain_with_cross_validation, which takes as input, a cross validation data module that has thetrain_val_dataloadersmethod, and a model config.NOTE: Currently, the model config is also apydanticmodel, but this requirement could be relaxed in the future.importlightningasLfromlightning_cvimportCrossValidationTrainerConfig,CrossValModuleclassCrossValidationTrainer:__fabric_keys__={"accelerator","strategy","devices","precision","plugins","callbacks","loggers",}def__init__(self,model_type:type[CrossValModule],config:CrossValidationTrainerConfig,):self.model_typeself.config=configself.fabric=L.Fabric(**self.config.model_dump(include=self.__fabric_keys__))# other init logicdeftrain_with_cross_validation(self,datamodule,model_config):# this method setups up k-fold models# then does the cross validation loop described in the# `Overall workflow` section aboveCallbacksPyTorch Lightninghas predefined hooks that enable custom logic to be applied at various steps of the train/inference loop. InLightning Fabric, no predefined hooks are automatically available, which provides authors flexibility of of defining what hooks are available.The following hooks are available when using theCrossValidationTrainer:HookCalledon_train_startonly oncewhen at the beginning oftrain_with_cross_validationon_train_endonly oncewhen at the end oftrain_with_cross_validationon_train_fold_startper epochbefore the fold level loopson_train_fold_endper epochafter the fold level loopson_train_epoch_start_per_foldper epoch per foldbefore batch level loopson_train_epoch_end_per_foldper epoch per foldafter batch level loopson_train_batch_start_per_foldper batchbefore forward/backward, logging, etcon_train_batch_end_per_foldper batchafter forward/backward, logging, etcon_before_optimizer_stepper batchbeforeoptimizer.stepon_before_zero_gradper batchbeforeoptimizer.zero_gradon_validation_start_per_foldper epoch per foldbefore batch level loopson_validation_end_per_foldper epoch per foldafter batch level loopson_before_log_metricsper batchbefore metrics are logged during BOTH train/val loopson_validation_batch_start_per_foldper batchbefore validation metrics are obtainedon_validation_batch_end_per_foldper batchafter validation metrics are obtainedon_before_backwardper batchbeforeloss.backwardon_after_backwardper batchafterloss.backwardA callback base class has been provided to to enable the hooks for any method. Thelightning_cv.callbacks.CallbackMUST be used to provide the hooks, since thepydantictrainer config will only accept subclasses of this exact class.An example callback to summarize model parameters that relies on alightningfunction:importsysfromtypingimportcastimportlightning_cvaslcvfromlightning.pytorch.utilities.model_summary.model_summaryimportsummarizeclassModelSummary(lcv.callbacks.Callback):def__init__(self,max_depth:int=1):self.max_depth=max_depth# we only need to generate the model summary once at the beginning of cross validationdefon_train_start(self,trainer:lcv.CrossValidationTrainer):model=trainer.fold_manager[0].modelsummary=summarize(model,max_depth=self.max_depth)sys.stderr.write(repr(summary))def__repr__(self)->str:returnf"{self.__class__.__name__}(max_depth={self.max_depth})"Then, any custom callbacks can be passed to thecallbacksargument of theCrossValidationTrainerConfig.Minimal cross validation referencefromlightning_cvimport(CrossValidationTrainer,CrossValidationTrainerConfig,CrossValModule,_CrossValDataModule)frompydanticimportBaseModelclassMyModelConfig:...classMyModel(CrossValModule):def__init__(self,config:MyModelConfig,...):super().__init__(...)...deftraining_step(self,...):...defvalidation_step(self,...):...defconfigure_optimizers(self,...):...classMyDataModule(_CrossValDataModule):def__init__(self,...):...deftrain_val_dataloaders(self,...):...# 1. setup datamoduledatamodule=MyDataModule(...)# 2. setup model init configmodel_config=MyModelConfig()# 3. setup trainerconfig=CrossValidationTrainerConfig()trainer=CrossValidationTrainer(model_type=MyModel,config=config)# 4. run cross validation#    handles setup logic for model, datamodule, and cv-foldstrainer.train_with_cross_validation(datamodule=datamodule,model_config=model_config)TuningCurrently, there is a simple integration foroptuna, a machine learning hyperparameter tuning library.Here is a simple example that uses the providedTunerclass to setup tuning runs. Suppose you want to tune the learning rate:importoptunafromlightning_cvimport(CrossValidationTrainer,CrossValidationTrainerConfig,CrossValModule,_CrossValDataModule)fromlightning_cv.tuningimportTunerfrompydanticimportBaseModelfromfunctoolsimportpartialclassMyModelConfig:lr:float...classMyModel(CrossValModule):def__init__(self,config:MyModelConfig,...):super().__init__(...)self.config=config...deftraining_step(self,...):...defvalidation_step(self,...):...defconfigure_optimizers(self,...):lr=self.config.lrclassMyDataModule(_CrossValDataModule):def__init__(self,...):...deftrain_val_dataloaders(self,...):...# 1. setup datamoduledatamodule=MyDataModule(...)# 2. setup model init configmodel_config=MyModelConfig()# 3. create a tunertuner=Tuner(model_type=MyModel,model_config=MyModelConfig(),datamodule=datamodule,trainer_config=CrossValidationTrainerConfig(),)# 4. create a tuning suggestion fn#    signature: (optuna.Trial) -> dict[str, Any]defget_trial_suggestions(trial:optuna.Trial):# the optuna trial object will sample a float# from the loguniform range [1e-4, 1e-2]lr=trial.suggest_float(name="lr",low=1e-4,high=1e-2,log=True)return{"lr":lr}# 5. run tuning trialsstudy=optuna.create_study(study_name="test",# look into other args to modify study such as# bad trial pruning)# we still don't have any actual trial obj so need a partial fn# that will take these later, thus the trial MUST# be the first argdef_tuning_trial(trial:optuna.Trial,tuner:Tuner):returntuner.tune(trial=trial,func=get_trial_suggestions)tuning_trial=partial(_tuning_trial,tuner=tuner)study.optimize(tuning_trial,n_trials=100,direction="minimize"# ie min loss)The only public method theTunerclass has is thetunemethod. This method first adjusts the model config which new trialed values, which is only the learning rate in this example. All other default values are held constant.Then, the tuner sets up a trainer instance and calls thetrain_with_cross_validation_methodwith the datamodule. If the method finishes and is not pruned due to an unpromising trial, then the.tunemethod returns a monitored validation metric, such as the validation loss, averaged over all folds.This validation metric is reported to theoptuna.Studyobject, which then keeps track of the best run.