Package: torchfitter
torchfitteris a simple library to ease the training of PyTorch models. It
features a class calledTrainerthat includes the basic functionality to fit
models in a Keras-like style.Internally,torchfitterleverages the power ofaccelerateto handle the device management.The library also provides a callbacks API that can be used to interact with
the model during the training process, as well as a set of basic regularization
procedures.InstallationNormal userpipinstalltorchfitterThis library does not ship CUDA nor XLA. Follow theofficial PyTorch documentationfor
more information about how to install CUDA binaries.Developergitclonehttps://github.com/Xylambda/torchfitter.git
pipinstall-etorchfitter/.-rtorchfitter/requirements-dev.txtTestsTo run the tests you must install the library as adeveloper.cdtorchfitter/
pytest-vtests/FeaturesSupportedNot supportedPlannedBasic training loopxGradient ClippingxGradient AccumulationxMulti-device supportxRegularizationxIn-loop metrics supportxMixed precision trainingxCallbacks SystemxHyperparameter searchxWarm TrainingxxUsageAssume we already haveDataLoadersfor the train and validation sets.fromtorch.utils.dataimportDataLoadertrain_loader=DataLoader(...)val_loader=DataLoader(...)Then, create the optimizer and the loss criterion as usual. Pass them to the
trainer along the PyTorch model. You can also add a regularization procedure if
you need/want to do it. The same goes for callbacks: create the desired
callbacks and pass them to the trainer as a list.importtorch.nnasnnimporttorch.optimasoptimfromtorchfitter.trainerimportTrainerfromtorchfitter.callbacksimport(LoggerCallback,EarlyStopping,LearningRateScheduler,L1Regularization,)model=nn.Linear(in_features=1,out_features=1)criterion=nn.MSELoss()optimizer=optim.Adam(model.parameters())l1_reg=L1Regularization(regularization_rate=0.01,biases=False)# callbackslogger=LoggerCallback(update_step=50)early_stopping=EarlyStopping(patience=50,load_best=True,path='checkpoint.pt')scheduler=LearningRateScheduler(scheduler=optim.lr_scheduler.StepLR(optimizer,step_size=500,gamma=0.9))trainer=Trainer(model=model,criterion=criterion,optimizer=optimizer,mixed_precision="fp16",accumulate_iter=4,# accumulate gradient every 4 iterations,gradient_clipping='norm',gradient_clipping_kwrgs={'max_norm':1.0,'norm_type':2.0},callbacks=[l1_reg,scheduler,early_stopping,logger])history=trainer.fit(train_loader,val_loader,epochs=1000)Sincetorchfitterleverages the power ofaccelerate, the device management
will rely on the latter. You can pass your ownaccelerate.Acceleratorobject to fine tune its parameters:fromaccelerateimportAcceleratorfromtorchfitter.trainerimportTraineraccelerator=Accelerator(...)trainer=Trainer(accelerator=accelerator,**kwargs)CallbacksCallbacks allow you to interact with the model during the fitting process. They
provide with different methods that are called at different stages. To create a
callback simply extend the base class and fill the desired methods.importtorchfromtorchfitter.conventionsimportParamsDictfromtorchfitter.callbacks.baseimportCallbackclassModelCheckpoint(Callback):def__init__(self,path):super(ModelCheckpoint,self).__init__()self.path=pathdef__repr__(self)->str:return"ModelCheckpoint()"defon_epoch_end(self,params_dict):# get paramsaccelerator=params_dict[ParamsDict.ACCELERATOR]epoch=params_dict[ParamsDict.EPOCH_NUMBER]# ensure model is safe to save_model=params_dict[ParamsDict.MODEL]accelerator.wait_for_everyone()unwrapped_model=accelerator.unwrap_model(_model)# actual savingfname=self.path/f'model_epoch_{epoch}.pt'accelerator.save(unwrapped_model.state_dict(),fname)Each method receivesparams_dict, which is a dictionary object containing the
internal training parameters. You can see the pair key value of each parameter
of the conventions:>>>fromtorchfitter.conventionsimportParamsDict>>>[(x,getattr(ParamsDict,x))forxinParamsDict.__dict__ifnotx.startswith('__')]And you can also check the doc to understand the meaning of each one of the
parameters:>>>fromtorchfitter.conventionsimportParamsDict>>>print(ParamsDict.__doc__)NOTE:the callbacks design can be considered as a port from Keras design.I AM NOTthe author of this callback sysem design despite the fact that I
made some minor design changes. Find more in theCreditssection.FAQDo you know Pytorch-Lightning/FastAI?I know them and I thinkthey are awesome. This is a personal project though
I must say the trainer is reasonably well-equiped.Why is thevalidation loadernot optional?Because I think it enforces good ML practices that way.Why didn't you implement the optimization steps in the model object?It is certainly another approach you may take when building an optimization
loop (PyTorch-Lightning works this way), but I don't like my abstract data
types to track way too many things in addition to being torch.nn.Module types.
Functionality should beclear and atomic: the model tracks gradients and
the trainer cares about the optimization process.I have a suggestion/questionThank you! Do not hesitate to open an issue and I'll do my best to answer you.CREDITSKeras API.torchmetricsfastaiCiteIf you've used this library for your projects please cite it:@misc{alejandro2019torchfitter,
  title={torchfitter - Simple Trainer to Optimize PyTorch Models},
  author={Alejandro Pérez-Sanjuán},
  year={2020},
  howpublished={\url{https://github.com/Xylambda/torchfitter}},}