2023-05-01 20:20:05 +00:00
---
comments: true
2023-07-17 10:40:04 +00:00
description: Discover how to customize and extend base Ultralytics YOLO Trainer engines. Support your custom model and dataloader by overriding built-in functions.
keywords: Ultralytics, YOLO, trainer engines, BaseTrainer, DetectionTrainer, customizing trainers, extending trainers, custom model, custom dataloader
2023-05-01 20:20:05 +00:00
---
2023-01-12 16:09:26 +00:00
Both the Ultralytics YOLO command-line and python interfaces are simply a high-level abstraction on the base engine
executors. Let's take a look at the Trainer engine.
2023-01-08 13:01:22 +00:00
## BaseTrainer
2023-02-04 15:54:34 +00:00
BaseTrainer contains the generic boilerplate training routine. It can be customized for any task based over overriding
2023-01-12 16:09:26 +00:00
the required functions or operations as long the as correct formats are followed. For example, you can support your own
2023-02-04 15:54:34 +00:00
custom model and dataloader by just overriding these functions:
2023-01-12 16:09:26 +00:00
* `get_model(cfg, weights)` - The function that builds the model to be trained
2023-02-04 15:54:34 +00:00
* `get_dataloder()` - The function that builds the dataloader
2023-07-16 15:47:45 +00:00
More details and source code can be found in [`BaseTrainer` Reference ](../reference/engine/trainer.md )
2023-01-08 13:01:22 +00:00
## DetectionTrainer
2023-01-12 16:09:26 +00:00
2023-01-08 13:01:22 +00:00
Here's how you can use the YOLOv8 `DetectionTrainer` and customize it.
2023-01-12 16:09:26 +00:00
2023-01-08 13:01:22 +00:00
```python
2023-07-16 15:47:45 +00:00
from ultralytics.models.yolo.detect import DetectionTrainer
2023-01-08 13:01:22 +00:00
trainer = DetectionTrainer(overrides={...})
trainer.train()
2023-01-12 16:09:26 +00:00
trained_model = trainer.best # get best model
2023-01-08 13:01:22 +00:00
```
### Customizing the DetectionTrainer
2023-01-12 16:09:26 +00:00
Let's customize the trainer **to train a custom detection model** that is not supported directly. You can do this by
simply overloading the existing the `get_model` functionality:
2023-01-08 13:01:22 +00:00
```python
2023-07-16 15:47:45 +00:00
from ultralytics.models.yolo.detect import DetectionTrainer
2023-01-12 16:09:26 +00:00
2023-01-08 13:01:22 +00:00
class CustomTrainer(DetectionTrainer):
def get_model(self, cfg, weights):
...
2023-01-12 16:09:26 +00:00
2023-01-08 13:01:22 +00:00
trainer = CustomTrainer(overrides={...})
trainer.train()
```
2023-01-12 16:09:26 +00:00
2023-01-08 13:01:22 +00:00
You now realize that you need to customize the trainer further to:
2023-05-25 16:12:50 +00:00
* Customize the `loss function` .
2023-01-12 16:09:26 +00:00
* Add `callback` that uploads model to your Google Drive after every 10 `epochs`
Here's how you can do it:
2023-01-08 13:01:22 +00:00
```python
2023-07-16 15:47:45 +00:00
from ultralytics.models.yolo.detect import DetectionTrainer
2023-05-25 16:12:50 +00:00
from ultralytics.nn.tasks import DetectionModel
2023-05-25 10:37:54 +00:00
class MyCustomModel(DetectionModel):
2023-05-25 16:12:50 +00:00
def init_criterion(self):
...
2023-01-12 16:09:26 +00:00
2023-01-08 13:01:22 +00:00
class CustomTrainer(DetectionTrainer):
def get_model(self, cfg, weights):
2023-05-25 10:37:54 +00:00
return MyCustomModel(...)
2023-01-08 13:01:22 +00:00
2023-05-25 16:12:50 +00:00
2023-01-08 13:01:22 +00:00
# callback to upload model weights
def log_model(trainer):
last_weight_path = trainer.last
...
2023-01-12 16:09:26 +00:00
2023-01-08 13:01:22 +00:00
trainer = CustomTrainer(overrides={...})
2023-01-12 16:09:26 +00:00
trainer.add_callback("on_train_epoch_end", log_model) # Adds to existing callback
2023-01-08 13:01:22 +00:00
trainer.train()
```
2023-04-05 22:25:32 +00:00
To know more about Callback triggering events and entry point, checkout our [Callbacks Guide ](callbacks.md )
2023-01-08 13:01:22 +00:00
## Other engine components
2023-01-12 16:09:26 +00:00
2023-02-04 15:54:34 +00:00
There are other components that can be customized similarly like `Validators` and `Predictors`
2023-06-15 19:17:10 +00:00
See Reference section for more information on these.