Reference for ultralytics/utils/torch_utils.py
Note
Full source code for this file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/torch_utils.py. Help us fix any issues you see by submitting a Pull Request 🛠️. Thank you 🙏!
ultralytics.utils.torch_utils.ModelEMA
Updated Exponential Moving Average (EMA) from https://github.com/rwightman/pytorch-image-models
Keeps a moving average of everything in the model state_dict (parameters and buffers)
For EMA details see https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage
To disable EMA set the enabled
attribute to False
.
Source code in ultralytics/utils/torch_utils.py
__init__(model, decay=0.9999, tau=2000, updates=0)
Create EMA.
Source code in ultralytics/utils/torch_utils.py
update(model)
Update EMA parameters.
Source code in ultralytics/utils/torch_utils.py
update_attr(model, include=(), exclude=('process_group', 'reducer'))
Updates attributes and saves stripped model with optimizer removed.
ultralytics.utils.torch_utils.EarlyStopping
Early stopping class that stops training when a specified number of epochs have passed without improvement.
Source code in ultralytics/utils/torch_utils.py
__call__(epoch, fitness)
Check whether to stop training
Parameters:
Name | Type | Description | Default |
---|---|---|---|
epoch |
int
|
Current epoch of training |
required |
fitness |
float
|
Fitness value of current epoch |
required |
Returns:
Type | Description |
---|---|
bool
|
True if training should stop, False otherwise |
Source code in ultralytics/utils/torch_utils.py
__init__(patience=50)
Initialize early stopping object
Parameters:
Name | Type | Description | Default |
---|---|---|---|
patience |
int
|
Number of epochs to wait after fitness stops improving before stopping. |
50
|
Source code in ultralytics/utils/torch_utils.py
ultralytics.utils.torch_utils.torch_distributed_zero_first(local_rank)
Decorator to make all processes in distributed training wait for each local_master to do something.
Source code in ultralytics/utils/torch_utils.py
ultralytics.utils.torch_utils.smart_inference_mode()
Applies torch.inference_mode() decorator if torch>=1.9.0 else torch.no_grad() decorator.
Source code in ultralytics/utils/torch_utils.py
ultralytics.utils.torch_utils.get_cpu_info()
Return a string with system CPU information, i.e. 'Apple M2'.
Source code in ultralytics/utils/torch_utils.py
ultralytics.utils.torch_utils.select_device(device='', batch=0, newline=False, verbose=True)
Selects the appropriate PyTorch device based on the provided arguments.
The function takes a string specifying the device or a torch.device object and returns a torch.device object representing the selected device. The function also validates the number of available devices and raises an exception if the requested device(s) are not available.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
device |
str | device
|
Device string or torch.device object. Options are 'None', 'cpu', or 'cuda', or '0' or '0,1,2,3'. Defaults to an empty string, which auto-selects the first available GPU, or CPU if no GPU is available. |
''
|
batch |
int
|
Batch size being used in your model. Defaults to 0. |
0
|
newline |
bool
|
If True, adds a newline at the end of the log string. Defaults to False. |
False
|
verbose |
bool
|
If True, logs the device information. Defaults to True. |
True
|
Returns:
Type | Description |
---|---|
device
|
Selected device. |
Raises:
Type | Description |
---|---|
ValueError
|
If the specified device is not available or if the batch size is not a multiple of the number of devices when using multiple GPUs. |
Examples:
Note
Sets the 'CUDA_VISIBLE_DEVICES' environment variable for specifying which GPUs to use.
Source code in ultralytics/utils/torch_utils.py
ultralytics.utils.torch_utils.time_sync()
ultralytics.utils.torch_utils.fuse_conv_and_bn(conv, bn)
Fuse Conv2d() and BatchNorm2d() layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/.
Source code in ultralytics/utils/torch_utils.py
ultralytics.utils.torch_utils.fuse_deconv_and_bn(deconv, bn)
Fuse ConvTranspose2d() and BatchNorm2d() layers.
Source code in ultralytics/utils/torch_utils.py
ultralytics.utils.torch_utils.model_info(model, detailed=False, verbose=True, imgsz=640)
Model information. imgsz may be int or list, i.e. imgsz=640 or imgsz=[640, 320].
Source code in ultralytics/utils/torch_utils.py
ultralytics.utils.torch_utils.get_num_params(model)
ultralytics.utils.torch_utils.get_num_gradients(model)
Return the total number of parameters with gradients in a YOLO model.
ultralytics.utils.torch_utils.model_info_for_loggers(trainer)
Return model info dict with useful model information.
Example
YOLOv8n info for loggers
Source code in ultralytics/utils/torch_utils.py
ultralytics.utils.torch_utils.get_flops(model, imgsz=640)
Return a YOLO model's FLOPs.
Source code in ultralytics/utils/torch_utils.py
ultralytics.utils.torch_utils.get_flops_with_torch_profiler(model, imgsz=640)
Compute model FLOPs (thop alternative).
Source code in ultralytics/utils/torch_utils.py
ultralytics.utils.torch_utils.initialize_weights(model)
Initialize model weights to random values.
Source code in ultralytics/utils/torch_utils.py
ultralytics.utils.torch_utils.scale_img(img, ratio=1.0, same_shape=False, gs=32)
Source code in ultralytics/utils/torch_utils.py
ultralytics.utils.torch_utils.make_divisible(x, divisor)
Returns nearest x divisible by divisor.
ultralytics.utils.torch_utils.copy_attr(a, b, include=(), exclude=())
Copies attributes from object 'b' to object 'a', with options to include/exclude certain attributes.
Source code in ultralytics/utils/torch_utils.py
ultralytics.utils.torch_utils.get_latest_opset()
Return second-most (for maturity) recently supported ONNX opset by this version of torch.
ultralytics.utils.torch_utils.intersect_dicts(da, db, exclude=())
Returns a dictionary of intersecting keys with matching shapes, excluding 'exclude' keys, using da values.
Source code in ultralytics/utils/torch_utils.py
ultralytics.utils.torch_utils.is_parallel(model)
ultralytics.utils.torch_utils.de_parallel(model)
De-parallelize a model: returns single-GPU model if model is of type DP or DDP.
ultralytics.utils.torch_utils.one_cycle(y1=0.0, y2=1.0, steps=100)
Returns a lambda function for sinusoidal ramp from y1 to y2 https://arxiv.org/pdf/1812.01187.pdf.
ultralytics.utils.torch_utils.init_seeds(seed=0, deterministic=False)
Initialize random number generator (RNG) seeds https://pytorch.org/docs/stable/notes/randomness.html.
Source code in ultralytics/utils/torch_utils.py
ultralytics.utils.torch_utils.strip_optimizer(f='best.pt', s='')
Strip optimizer from 'f' to finalize training, optionally save as 's'.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
f |
str
|
file path to model to strip the optimizer from. Default is 'best.pt'. |
'best.pt'
|
s |
str
|
file path to save the model with stripped optimizer to. If not provided, 'f' will be overwritten. |
''
|
Returns:
Type | Description |
---|---|
None
|
None |
Example
Source code in ultralytics/utils/torch_utils.py
ultralytics.utils.torch_utils.profile(input, ops, n=10, device=None)
Ultralytics speed, memory and FLOPs profiler.