Train & Validation

Training Model

To train a model, the ModelTrainer can help manage the training process. Initialize the ModelTrainer and use the solve() function to start the training.

Before starting the training, don’t forget to start the progress logger to enable logging the process status. This will also enable Weights & Biases (wandb) or TensorBoard if configured.

from yolo import ModelTrainer
solver = ModelTrainer(cfg, model, converter, progress, device, use_ddp)
progress.start()
solver.solve(dataloader)

Training Diagram

The following diagram illustrates the training process:

flowchart LR subgraph TS["trainer.solve"] subgraph TE["train one epoch"] subgraph "train one batch" backpropagation-->TF[forward] TF-->backpropagation end end subgraph validator.solve VC["calculate mAP"]-->VF[forward] VF[forward]-->VC end end TE-->validator.solve validator.solve-->TE

Validation Model

To validate the model performance, we follow a similar approach as the training process using ModelValidator.

from yolo import ModelValidator
solver = ModelValidator(cfg, model, converter, progress, device, use_ddp)
progress.start()
solver.solve(dataloader)

The ModelValidator class helps manage the validation process, ensuring that the model’s performance is evaluated accurately.

Note

The original training process already includes the validation phase. Call this separately if you want to run the validation again after the training is completed.