Reference for ultralytics/utils/torch_utils.py
Note
This file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/torch_utils.py. If you spot a problem please help fix it by contributing 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
78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 |
|
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)
Scales and pads an image tensor of shape img(bs,3,y,x) based on given ratio and grid size gs, optionally retaining the original shape.
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.convert_optimizer_state_dict_to_fp16(state_dict)
Converts the state_dict of a given optimizer to FP16, focusing on the 'state' key for tensor conversions.
This method aims to reduce storage size without altering 'param_groups' as they contain non-tensor data.
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.