Package: torchmanager
torchmanagerA generic deep learning training/testing framework for PyTorchTo use this framework, simply initialize aManagerobject. TheManagerclass provides a generic training/testing loop for PyTorch models. It also provides some useful callbacks to use during training/testing.Pre-requestThe following packages are required to use this framework:Python 3.8+PyTorchtqdmscipy (Optional for FID metric)tensorboard (Optional for tensorboard recording)InstallationThe package can be installed directly from PyPi or Conda:PyPi:pip install torchmanagerConda:conda install -c kisonho torchmanagerStart from ConfigurationsTheConfigsclass is designed to be inherited to define necessary configurations. It also provides a method to get configurations from terminal arguments.fromtorchmanager.configsimportConfigsas_Configs# define necessary configurationsclassConfigs(_Configs):epochs:intlr:float...defget_arguments(parser:Union[argparse.ArgumentParser,argparse._ArgumentGroup]=argparse.ArgumentParser())->Union[argparse.ArgumentParser,argparse._ArgumentGroup]:'''Add arguments to argument parser'''...defshow_settings(self)->None:...# get configs from terminal argumentsconfigs=Configs.from_arguments()Torchmanager DatasetThedata.Datasetclass is designed to be inherited to define a dataset. It is a combination oftorch.utils.data.Datasetandtorch.utils.data.DataLoaderwith easier usage.fromtorchmanager.dataimportDataset# define datasetclassCustomDataset(Dataset):def__init__(self,...):...@propertydefunbatched_len(self)->int:...def__getitem__(self,index:int)->Tuple[torch.Tensor,torch.Tensor]:...# initialize datasetstraining_dataset=CustomDataset(...)val_dataset=CustomDataset(...)testing_dataset=CustomDataset(...)The ManagerTheManagerclass is the core of the framework. It provides a generic training/testing loop for PyTorch models. TheManagerclass is designed to be inherited to manage the training/testing algorithm. There are also some useful callbacks to use during training/testing.Initialize the manager with target model, optimizer, loss function, and metrics:importtorch,torchmanager# define modelclassPytorchModel(torch.nn.Module):...# initialize model, optimizer, loss function, and metricsmodel=PytorchModel(...)optimizer=torch.optim.SGD(model.parameters(),lr=configs.lr)loss_fn=torchmanager.losses.CrossEntropy()metrics={'accuracy':torchmanager.metrics.SparseCategoricalAccuracy()}# initialize managermanager=torchmanager.Manager(model,optimizer,loss_fn=loss_fn,metrics=metrics)Multiple losses can be used by passing a dictionary toloss_fn:loss_fn={'loss1':torchmanager.losses.CrossEntropy(),'loss2':torchmanager.losses.Dice(),...}Train the model with `fit`` method:show_verbose:bool=...# show progress bar information during training/testingmanager.fit(training_dataset,epochs=configs.epochs,val_dataset=val_dataset,show_verbose=show_verbose)There are also some other callbacks to use:tensorboard_callback=torchmanager.callbacks.TensorBoard('logs')# tensorboard dependency requiredlast_ckpt_callback=torchmanager.callbacks.LastCheckpoint(manager,'last.model')model=manager.fit(...,callbacks_list=[tensorboard_callback,last_ckpt_callback])Test the model with test method:manager.test(testing_dataset,show_verbose=show_verbose)Save the final trained PyTorch model:torch.save(model,"model.pth")# The saved PyTorch model can be loaded individually without using torchmanagerDevice selection during training/testingTorchmanager automatically detects available devices to use during training/testing. GPU/MPS will be used in first priority if available. To specify other device to use, simply pass the device to thefitmethod for training andtestmethod for testing:Multi-GPU training/testing:# train on multiple GPUsmodel=manager.fit(...,use_multi_gpus=True)# test on multiple GPUsmanager.test(...,use_multi_gpus=True)Use only specified GPUs for training/testing:# specify devices to usegpus:Union[list[torch.device],torch.device]=...# Notice: device id must be specified# train on specified multiple GPUsmodel=manager.fit(...,use_multi_gpus=True,devices=gpus)# test on specified multiple GPUsmanager.test(...,use_multi_gpus=True,devices=gpus)Customize training/testing algorithmTheManagerclass is designed to be inherited to manage the training/testing algorithm. To customize the training/testing algorithm, simply inherit theManagerclass and override thetrain_stepandtest_stepmethods.classCustomManager(Manager):...deftrain_step(x_train:torch.Tensor,y_train:torch.Tensor)->Dict[str,float]:...deftest_step(x_test:torch.Tensor,y_test:torch.Tensor)->Dict[str,float]:...The saved experiment informationTheExperimentclass is designed to be used as a single callback to save experiment information. It is a combination oftorchmanager.callbacks.TensorBoard,torchmanager.callbacks.LastCheckpoint, andtorchmanager.callbacks.BestCheckpointwith easier usage....exp_callback=torchmanager.callbacks.Experiment('test.exp',manager)# tensorboard dependency requiredmodel=manager.fit(...,callbacks_list=[exp_callback])The information, including full training logs and checkpoints, will be saved in the following structure:experiments
└── <experiment name>.exp
    ├── checkpoints
    │   ├── best-<metric name>.model
    │   └── last.model
    └── data
    │   └── <TensorBoard data file>
    └── <experiment name>.logPlease cite this work if you find it useful@software{he_2023_10381715,author={He, Qisheng andDong, Ming},title={{TorchManager: A generic deep learningtraining/testing framework for PyTorch}},month=dec,year=2023,publisher={Zenodo},version=1,doi={10.5281/zenodo.10381715},url={https://doi.org/10.5281/zenodo.10381715}}Also checkout our projects implemented with torchmanagermagnet- Modality-Agnostic Learning for Medical Image Segmentation Using Multi-modality Self-distillationtlt- Transferring Lottery Tickets in Computer Vision Models: a Dynamic Pruning Approach