You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
95 lines
3.6 KiB
95 lines
3.6 KiB
--- |
|
comments: true |
|
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 |
|
--- |
|
|
|
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. |
|
|
|
<p align="center"> |
|
<br> |
|
<iframe loading="lazy" width="720" height="405" src="https://www.youtube.com/embed/GsXGnb-A4Kc?start=104" |
|
title="YouTube video player" frameborder="0" |
|
allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture; web-share" |
|
allowfullscreen> |
|
</iframe> |
|
<br> |
|
<strong>Watch:</strong> Mastering Ultralytics YOLOv8: Advanced Customization |
|
</p> |
|
|
|
## BaseTrainer |
|
|
|
BaseTrainer contains the generic boilerplate training routine. It can be customized for any task based over overriding the required functions or operations as long the as correct formats are followed. For example, you can support your own custom model and dataloader by just overriding these functions: |
|
|
|
- `get_model(cfg, weights)` - The function that builds the model to be trained |
|
- `get_dataloader()` - The function that builds the dataloader More details and source code can be found in [`BaseTrainer` Reference](../reference/engine/trainer.md) |
|
|
|
## DetectionTrainer |
|
|
|
Here's how you can use the YOLOv8 `DetectionTrainer` and customize it. |
|
|
|
```python |
|
from ultralytics.models.yolo.detect import DetectionTrainer |
|
|
|
trainer = DetectionTrainer(overrides={...}) |
|
trainer.train() |
|
trained_model = trainer.best # get best model |
|
``` |
|
|
|
### Customizing the DetectionTrainer |
|
|
|
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: |
|
|
|
```python |
|
from ultralytics.models.yolo.detect import DetectionTrainer |
|
|
|
|
|
class CustomTrainer(DetectionTrainer): |
|
def get_model(self, cfg, weights): |
|
"""Loads a custom detection model given configuration and weight files.""" |
|
... |
|
|
|
|
|
trainer = CustomTrainer(overrides={...}) |
|
trainer.train() |
|
``` |
|
|
|
You now realize that you need to customize the trainer further to: |
|
|
|
- Customize the `loss function`. |
|
- Add `callback` that uploads model to your Google Drive after every 10 `epochs` Here's how you can do it: |
|
|
|
```python |
|
from ultralytics.models.yolo.detect import DetectionTrainer |
|
from ultralytics.nn.tasks import DetectionModel |
|
|
|
|
|
class MyCustomModel(DetectionModel): |
|
def init_criterion(self): |
|
"""Initializes the loss function and adds a callback for uploading the model to Google Drive every 10 epochs.""" |
|
... |
|
|
|
|
|
class CustomTrainer(DetectionTrainer): |
|
def get_model(self, cfg, weights): |
|
"""Returns a customized detection model instance configured with specified config and weights.""" |
|
return MyCustomModel(...) |
|
|
|
|
|
# callback to upload model weights |
|
def log_model(trainer): |
|
"""Logs the path of the last model weight used by the trainer.""" |
|
last_weight_path = trainer.last |
|
print(last_weight_path) |
|
|
|
|
|
trainer = CustomTrainer(overrides={...}) |
|
trainer.add_callback("on_train_epoch_end", log_model) # Adds to existing callback |
|
trainer.train() |
|
``` |
|
|
|
To know more about Callback triggering events and entry point, checkout our [Callbacks Guide](callbacks.md) |
|
|
|
## Other engine components |
|
|
|
There are other components that can be customized similarly like `Validators` and `Predictors`. See Reference section for more information on these.
|
|
|