Package: trainer
👟 TrainerAn opinionated general purpose model trainer on PyTorch with a simple code base.InstallationFrom Github:git clone https://github.com/coqui-ai/Trainercd Trainermake installFrom PyPI:pip install trainerPrefer installing from Github as it is more stable.Implementing a modelSubclass and overload the functions in theTrainerModel()Training a model with auto-optimizationSee theMNIST example.Training a model with advanced optimizationWith 👟 you can define the whole optimization cycle as you want as the in GAN example below. It enables more
under-the-hood control and flexibility for more advanced training loops.You just have to use thescaled_backward()function to handle mixed precision training....defoptimize(self,batch,trainer):imgs,_=batch# sample noisez=torch.randn(imgs.shape[0],100)z=z.type_as(imgs)# train discriminatorimgs_gen=self.generator(z)logits=self.discriminator(imgs_gen.detach())fake=torch.zeros(imgs.size(0),1)fake=fake.type_as(imgs)loss_fake=trainer.criterion(logits,fake)valid=torch.ones(imgs.size(0),1)valid=valid.type_as(imgs)logits=self.discriminator(imgs)loss_real=trainer.criterion(logits,valid)loss_disc=(loss_real+loss_fake)/2# step dicriminator_,_=self.scaled_backward(loss_disc,None,trainer,trainer.optimizer[0])iftrainer.total_steps_done%trainer.grad_accum_steps==0:trainer.optimizer[0].step()trainer.optimizer[0].zero_grad()# train generatorimgs_gen=self.generator(z)valid=torch.ones(imgs.size(0),1)valid=valid.type_as(imgs)logits=self.discriminator(imgs_gen)loss_gen=trainer.criterion(logits,valid)# step generator_,_=self.scaled_backward(loss_gen,None,trainer,trainer.optimizer[1])iftrainer.total_steps_done%trainer.grad_accum_steps==0:trainer.optimizer[1].step()trainer.optimizer[1].zero_grad()return{"model_outputs":logits},{"loss_gen":loss_gen,"loss_disc":loss_disc}...See theGAN training examplewith Gradient AccumulationTraining with Batch Size Findersee the test scriptherefor training with batch size finder.The batch size finder starts at a default BS(defaults to 2048 but can also be user defined) and searches for the largest batch size that can fit on your hardware. you should expect for it to run multiple trainings until it finds it. to use it instead of callingtrainer.fit()youll calltrainer.fit_with_largest_batch_size(starting_batch_size=2048)withstarting_batch_sizebeing the batch the size you want to start the search with. very useful if you are wanting to use as much gpu mem as possible.Training with DDP$python-mtrainer.distribute--scriptpath/to/your/train.py--gpus"0,1"We don't use.spawn()to initiate multi-gpu training since it causes certain limitations.Everything must the pickable..spawn()trains the model in subprocesses and the model in the main process is not updated.DataLoader with N processes gets really slow when the N is large.Training withAccelerateSettinguse_accelerateinTrainingArgstoTruewill enable training with Accelerate.You can also use it for multi-gpu or distributed training.CUDA_VISIBLE_DEVICES="0,1,2" accelerate launch --multi_gpu --num_processes 3 train_recipe_autoregressive_prompt.pySee theAccelerate docs.Adding a callback👟 Supports callbacks to customize your runs. You can either set callbacks in your model implementations or give them
explicitly to the Trainer.Please checktrainer.utils.callbacksto see available callbacks.Here is how you provide an explicit call back to a 👟Trainer object for weight reinitialization.defmy_callback(trainer):print(" > My callback was called.")trainer=Trainer(...,callbacks={"on_init_end":my_callback})trainer.fit()Profiling exampleCreate the torch profiler as you like and pass it to the trainer.importtorchprofiler=torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU,torch.profiler.ProfilerActivity.CUDA,],schedule=torch.profiler.schedule(wait=1,warmup=1,active=3,repeat=2),on_trace_ready=torch.profiler.tensorboard_trace_handler("./profiler/"),record_shapes=True,profile_memory=True,with_stack=True,)prof=trainer.profile_fit(profiler,epochs=1,small_run=64)thenrunTensorboardRun the tensorboard.tensorboard --logdir="./profiler/"Supported Experiment LoggersTensorboard- actively maintainedClearML- actively maintainedMLFlowAimWandDBTo add a new logger, you must subclassBaseDashboardLoggerand overload its functions.Anonymized TelemetryWe constantly seek to improve 🐸 for the community. To understand the community's needs better and address them accordingly, we collect stripped-down anonymized usage stats when you run the trainer.Of course, if you don't want, you can opt out by setting the environment variableTRAINER_TELEMETRY=0.