Package: wah
Installpip install wahRequirementsYou might want to manually installPyTorchfor GPU computation.lightning
PyYAML
tensorboard
torch
torchaudio
torchmetrics
torchvisionModel TrainingLet's trainResNet50[1]
onCIFAR-10[2] dataset
(Full example codeshere).First, import the package.importwahSecond, write your ownconfig.yamlfile (which will doeverythingfor you).num_classes:10batch_size:128num_workers:2epochs:200init_lr:0.1seed:0gpu:[0,]optimizer:SGDoptimizer_cfg:momentum:0.9weight_decay:5.e-4lr_scheduler:MultiStepLRlr_scheduler_cfg:milestones:[60,120,160,]gamma:0.2warmup_lr_scheduler:LinearLRwarmup_lr_scheduler_cfg:start_factor:0.01total_iters:10criterion:CrossEntropyLossmetrics:acc@1:Accuracy:task:multiclasstop_k:1ece:CalibrationError:task:multiclassnum_classes(int) -
number of classes in the dataset.batch_size(int) -
how many samples per batch to load.num_workers(int) -
how many subprocesses to use for data loading.
0 means that the data will be loaded in the main process.epochs(int) -
stop training once this number of epochs is reached.init_lr(float) -
initial learning rate.seed(int) -
seed value for random number generation.
Must be a non-negative integer.
If a negative integer is provided, no seeding will occur.gpu(List[int], optional) -
the GPU device(s) to be utilized for computation.
If not specified, the system automatically detects and selects devices.optimizer(str) -
specifies which optimizer to use.
Must be one of the optimizers supported intorch.optim.optimizer_cfg-
parameters for the specified optimizer.
Theparamsandlrparameters do not need to be explicitly provided (automatically initialized).lr_scheduler(str) -
specifies which scheduler to use.
Must be one of the schedulers supported intorch.optim.lr_scheduler.lr_scheduler_cfg-
parameters for the specified scheduler.
Theoptimizerparameter does not need to be explicitly provided (automatically initialized).warmup_lr_scheduler(str, optional) -
specifies which scheduler to use for warmup phase.
Must be one of ["ConstantLR","LinearLR", ].warmup_lr_scheduler_cfg(optional) -
parameters for the specified warmup scheduler.
Theoptimizerparameter does not need to be explicitly provided (automatically initialized).
Note that thetotal_itersparameter initializes the length of warmup phase.criterion(str) -
specifies which loss function to use.
Must be one of the loss functions supported intorch.nn.metrics-
metrics to record during the training and validation stages.
Must be the metrics supported intorchmetrics.Third, load yourconfig.yamlfile.config=wah.load_config(PATH_TO_CONFIG)Fourth, load your dataloaders.fromtorchvision.datasetsimportCIFAR10train_dataset=CIFAR10(train=True,...)val_dataset=CIFAR10(train=False,...)train_dataloader=wah.load_dataloader(dataset=train_dataset,config=config,shuffle=True,)val_dataloader=wah.load_dataloader(dataset=val_dataset,config=config,shuffle=False,)Fifth, load your model.fromtorchvision.modelsimportresnet50model=resnet50(weights=None,num_classes=10)model=wah.Wrapper(model,config)Finally, train your model!trainer=wah.load_trainer(config=config,save_dir=TRAIN_LOG_ROOT,name="cifar10-resnet50",every_n_epochs=SAVE_CKPT_PER_THIS_EPOCH,)trainer.fit(model=model,train_dataloaders=train_dataloader,val_dataloaders=val_dataloader,)You can check your train logs by running the following command:tensorboard --logdir TRAIN_LOG_ROOTReferences[1] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep Residual Learning for Image Recognition. CVPR, 2016.[2] Alex Krizhevsky and Geoffrey Hinton. Learning Multiple Layers of Features from Tiny Images. Tech. Rep., University
of Toronto, Toronto, Ontario, 2009.