Reference for ultralytics/models/yolo/detect/train.py
Note
This file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/models/yolo/detect/train.py. If you spot a problem please help fix it by contributing a Pull Request 🛠️. Thank you 🙏!
ultralytics.models.yolo.detect.train.DetectionTrainer
Bases: BaseTrainer
A class extending the BaseTrainer class for training based on a detection model.
Example
Parameters:
Name | Type | Description | Default |
---|---|---|---|
cfg | str | Path to a configuration file. Defaults to DEFAULT_CFG. | DEFAULT_CFG |
overrides | dict | Configuration overrides. Defaults to None. | None |
Source code in ultralytics/engine/trainer.py
auto_batch
Get batch size by calculating memory occupation of model.
Source code in ultralytics/models/yolo/detect/train.py
build_dataset
Build YOLO Dataset.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
img_path | str | Path to the folder containing images. | required |
mode | str |
| 'train' |
batch | int | Size of batches, this is for | None |
Source code in ultralytics/models/yolo/detect/train.py
get_dataloader
Construct and return dataloader.
Source code in ultralytics/models/yolo/detect/train.py
get_model
Return a YOLO detection model.
get_validator
Returns a DetectionValidator for YOLO model validation.
Source code in ultralytics/models/yolo/detect/train.py
label_loss_items
Returns a loss dict with labelled training loss items tensor.
Not needed for classification but necessary for segmentation & detection
Source code in ultralytics/models/yolo/detect/train.py
plot_metrics
plot_training_labels
Create a labeled training plot of the YOLO model.
Source code in ultralytics/models/yolo/detect/train.py
plot_training_samples
Plots training samples with their annotations.
Source code in ultralytics/models/yolo/detect/train.py
preprocess_batch
Preprocesses a batch of images by scaling and converting to float.
Source code in ultralytics/models/yolo/detect/train.py
progress_string
Returns a formatted string of training progress with epoch, GPU memory, loss, instances and size.
Source code in ultralytics/models/yolo/detect/train.py
set_model_attributes
Nl = de_parallel(self.model).model[-1].nl # number of detection layers (to scale hyps).