Reference for ultralytics/utils/autobatch.py
Note
This file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/autobatch.py. If you spot a problem please help fix it by contributing a Pull Request 🛠️. Thank you 🙏!
ultralytics.utils.autobatch.check_train_batch_size
Compute optimal YOLO training batch size using the autobatch() function.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
model | Module | YOLO model to check batch size for. | required |
imgsz | int | Image size used for training. | 640 |
amp | bool | Use automatic mixed precision if True. | True |
batch | float | Fraction of GPU memory to use. If -1, use default. | -1 |
Returns:
Type | Description |
---|---|
int | Optimal batch size computed using the autobatch() function. |
Note
If 0.0 < batch < 1.0, it's used as the fraction of GPU memory to use. Otherwise, a default fraction of 0.6 is used.
Source code in ultralytics/utils/autobatch.py
ultralytics.utils.autobatch.autobatch
Automatically estimate the best YOLO batch size to use a fraction of the available CUDA memory.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
model | module | YOLO model to compute batch size for. | required |
imgsz | int | The image size used as input for the YOLO model. Defaults to 640. | 640 |
fraction | float | The fraction of available CUDA memory to use. Defaults to 0.60. | 0.6 |
batch_size | int | The default batch size to use if an error is detected. Defaults to 16. | batch |
Returns:
Type | Description |
---|---|
int | The optimal batch size. |