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. |