ultralytics\utils\autobatch.py
# Ultralytics YOLO ?, AGPL-3.0 license
"""Functions for estimating the best YOLO batch size to use a fraction of the available CUDA memory in PyTorch."""
# 用于估算最佳 YOLO 批量大小的函数，以使用 PyTorch 中可用 CUDA 内存的一小部分。
from copy import deepcopy
import numpy as np
import torch
from ultralytics.utils import DEFAULT_CFG, LOGGER, colorstr
from ultralytics.utils.torch_utils import autocast, profile
# 这段代码定义了一个名为 check_train_batch_size 的函数，它用于确定在给定图像尺寸和内存使用率下，模型训练时的最佳批量大小。
# 1.model ：要评估的模型。
# 2.imgsz ：输入图像的尺寸，默认为 640。
# 3.amp ：布尔值，指示是否使用自动混合精度（Automatic Mixed Precision, AMP）。
# 4.batch ：批量大小的分数或具体数值。如果这个值在 0.0 到 1.0 之间，它表示 CUDA 内存使用率的分数；如果是一个具体的数值，则直接表示批量大小。
def check_train_batch_size(model, imgsz=640, amp=True, batch=-1):
# 使用 autobatch() 函数计算最佳 YOLO 训练批次大小。
"""
Compute optimal YOLO training batch size using the autobatch() function.
Args:
model (torch.nn.Module): YOLO model to check batch size for.
imgsz (int, optional): Image size used for training.
amp (bool, optional): Use automatic mixed precision if True.
batch (float, optional): Fraction of GPU memory to use. If -1, use default.
Returns:
(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.
"""
# 自动混合精度上下文管理。 使用 PyTorch 的 autocast 上下文管理器来启用或禁用自动混合精度。这有助于在训练过程中减少内存使用和加速计算，同时保持模型的精度。
# def autocast(enabled: bool, device: str = "cuda"):
# -> 它用于创建一个上下文管理器，以便在 PyTorch 的自动混合精度（AMP）中使用。如果 PyTorch 版本是 1.13 或更高，使用 torch.amp.autocast 并指定设备类型和启用状态。如果 PyTorch 版本低于 1.13，使用 torch.cuda.amp.autocast ，这时不能指定设备类型，只能启用或禁用自动混合精度。
# -> return torch.amp.autocast(device, enabled=enabled) / return torch.cuda.amp.autocast(enabled)
with autocast(enabled=amp):
# 调用 autobatch 函数。
# 使用 deepcopy 函数创建模型的一个深拷贝，以确保不会修改原始模型。
# 将模型设置为训练模式。
# 调用 autobatch 函数来计算最佳批量大小。 autobatch 函数接受 模型 、 图像尺寸 和 内存使用率的分数 作为参数。
# 如果 batch 参数在 0.0 到 1.0 之间，它被用作内存使用率的分数；否则，默认使用 0.6 作为内存使用率的分数。
# 返回值。函数返回 autobatch 函数计算出的最佳批量大小。
return autobatch(deepcopy(model).train(), imgsz, fraction=batch if 0.0 < batch < 1.0 else 0.6)
# 这个方法的目的是自动调整模型的批量大小，以便在不超过指定 CUDA 内存使用率的情况下最大化模型性能。通过这种方式，可以确保模型在有限的硬件资源下高效运行，同时避免内存溢出错误。# 这段代码定义了一个名为 autobatch 的函数，其目的是自动计算并优化模型在给定 CUDA 内存使用率下的最佳批量大小（batch size）。
# 1.model ：要评估的模型。
# 2.imgsz ：输入图像的尺寸，默认为 640。
# 3.fraction ：CUDA 内存使用率的分数，默认为 0.60（即 60%）。
# 4.batch_size ：默认的批量大小，通常从配置中获取。
def autobatch(model, imgsz=640, fraction=0.60, batch_size=DEFAULT_CFG.batch):
# 自动估计最佳 YOLO 批次大小以使用可用 CUDA 内存的一小部分。
"""
Automatically estimate the best YOLO batch size to use a fraction of the available CUDA memory.
Args:
model (torch.nn.module): YOLO model to compute batch size for.
imgsz (int, optional): The image size used as input for the YOLO model. Defaults to 640.
fraction (float, optional): The fraction of available CUDA memory to use. Defaults to 0.60.
batch_size (int, optional): The default batch size to use if an error is detected. Defaults to 16.
Returns:
(int): The optimal batch size.
"""
# Check device
# 这段代码是 autobatch 函数的开头部分，它负责检查模型所在的设备是否适合执行自动批量大小调整，并验证是否满足必要的条件。
# 定义前缀。 使用 colorstr 函数来定义一个前缀，用于所有日志消息，以便于识别。
# def colorstr(*input):
# -> 它用于生成带有ANSI转义序列的字符串，这些转义序列可以使终端中的文本显示不同的颜色和样式。
# -> 函数通过遍历 args 中的每个元素（颜色或样式），从 colors 字典中获取对应的ANSI转义序列，并将其与传入的 string 字符串连接起来。最后，它还会添加一个 colors["end"] 序列，用于重置终端的颜色和样式到默认状态。
# -> return "".join(colors[x] for x in args) + f"{string}" + colors["end"]
prefix = colorstr("AutoBatch: ")
# 记录信息。使用日志记录器 LOGGER 记录一条信息，表明正在计算给定图像尺寸 imgsz 和 CUDA 内存使用率 fraction 下的最佳批量大小。
LOGGER.info(f"{prefix}Computing optimal batch size for imgsz={imgsz} at {fraction * 100}% CUDA memory utilization.") # {prefix}在 CUDA 内存利用率为 {fraction * 100}% 时，计算 imgsz={imgsz} 的最佳批量大小。
# 获取模型设备。通过访问模型参数迭代器中的下一个元素来获取模型所在的设备（例如 GPU 或 CPU）。
device = next(model.parameters()).device # get model device
# 检查设备类型。 检查设备是否为 CPU 或 Apple Metal Performance Shaders (MPS)。如果是，说明代码不适用于这些设备，因为 autobatch 函数旨在用于 CUDA 设备（GPU）。
if device.type in {"cpu", "mps"}:
# 记录警告并返回默认批量大小。如果设备不是 CUDA 设备，记录一条警告信息，指出该函数仅适用于 CUDA 设备，并使用默认的批量大小。
LOGGER.info(f"{prefix} ⚠️ intended for CUDA devices, using default batch-size {batch_size}") # {prefix} ⚠️适用于 CUDA 设备，使用默认批量大小 {batch_size}。
# 返回默认的批量大小，并结束函数执行。
return batch_size
# 检查 torch.backends.cudnn.benchmark 设置。检查 PyTorch 的 cudnn.benchmark 设置是否为 True 。这个设置会影响 PyTorch 如何选择算法和卷积实现，可能会影响内存使用和性能。
if torch.backends.cudnn.benchmark:
LOGGER.info(f"{prefix} ⚠️ Requires torch.backends.cudnn.benchmark=False, using default batch-size {batch_size}") # {prefix} ⚠️ 需要 torch.backends.cudnn.benchmark=False，使用默认批量大小 {batch_size}。
# 记录警告并返回默认批量大小。如果 cudnn.benchmark 设置为 True ，记录一条警告信息，指出需要将其设置为 False 才能正确执行自动批量大小调整，并使用默认的批量大小。返回默认的批量大小，并结束函数执行。
return batch_size
# 这段代码的目的是确保 autobatch 函数在正确的环境下执行，并满足所有必要的条件。如果条件不满足，函数将记录适当的警告信息，并返回默认的批量大小。
# Inspect CUDA memory 检查 CUDA 内存。
# 这段代码是 autobatch 函数中用于检查和记录 CUDA 设备内存状态的部分。
# 定义字节到吉字节的转换。将字节转换为吉字节（GiB），因为 1 GiB 等于 2^30 字节。
gb = 1 << 30 # bytes to GiB (1024 ** 3)
# 获取设备名称。将设备对象转换为字符串并转换为大写，通常格式为 'CUDA:0'，表示第一个 CUDA 设备。
d = str(device).upper() # 'CUDA:0'
# 获取设备属性。使用 PyTorch 的 torch.cuda.get_device_properties 函数获取 CUDA 设备的属性。
properties = torch.cuda.get_device_properties(device) # device properties
# 计算总内存、已保留内存、已分配内存和空闲内存。
# 计算 CUDA 设备的总内存（以吉字节为单位）。
t = properties.total_memory / gb # GiB total
# 计算 CUDA 设备已保留的内存（以吉字节为单位）。
r = torch.cuda.memory_reserved(device) / gb # GiB reserved
# 计算 CUDA 设备已分配的内存（以吉字节为单位）。
a = torch.cuda.memory_allocated(device) / gb # GiB allocated
# 计算 CUDA 设备的空闲内存（以吉字节为单位）。
f = t - (r + a) # GiB free
# 记录 CUDA 内存状态。使用日志记录器 LOGGER 记录 CUDA 设备的内存状态，包括设备名称、总内存、已保留内存、已分配内存和空闲内存。
LOGGER.info(f"{prefix}{d} ({properties.name}) {t:.2f}G total, {r:.2f}G reserved, {a:.2f}G allocated, {f:.2f}G free") # {prefix}{d} ({properties.name}) {t:.2f}G 总计, {r:.2f}G 保留, {a:.2f}G 分配, {f:.2f}G 空闲。
# Profile batch sizes 配置文件批次大小。
# 这段代码是 autobatch 函数的核心部分，它负责通过试验不同的批量大小并测量它们对 CUDA 内存的影响来确定最佳批量大小。
# 变量定义。一个包含不同批量大小的列表，用于测试。
batch_sizes = [1, 2, 4, 8, 16]
# 尝试执行。尝试执行以下操作，如果出现任何异常，则执行 except 块。
try:
# 创建空图像张量。为每个批量大小创建一个空的图像张量，形状为 (b, 3, imgsz, imgsz) 。
img = [torch.empty(b, 3, imgsz, imgsz) for b in batch_sizes]
# 评估每个批量大小。使用 profile 函数评估每个批量大小的内存使用情况。 n=3 可能表示每个批量大小被评估 3 次。
# def profile(input, ops, n=10, device=None): -> 它用于评估模型或操作（ops）在不同输入下的性能和资源消耗。返回包含所有性能评估结果的 results 列表。记录 参数数量 、 GFLOPs 、 内存使用 、 前向传播时间 、 反向传播时间 、 输入 和 输出的形状 。 -> return results
results = profile(img, model, n=3, device=device)
# Fit a solution 找到解决方案。
# 拟合解决方案。
# 从 results 中提取每个批量大小的内存使用情况（假设内存使用是 x[2] ）。
y = [x[2] for x in results if x] # memory [2]
# numpy.polyfit(x, y, deg, rcond=None, full=False, w=None, cov=False)
# np.polyfit 是 NumPy 库中的一个函数，用于拟合数据到一个多项式模型。
# 参数 ：
# x ：一个数组，包含自变量的值。
# y ：一个数组，包含因变量的值。
# deg ：一个整数，表示拟合多项式的阶数。
# rcond ：一个浮点数，用于确定多项式系数矩阵的条件数。如果为 None ，则使用机器精度。
# full ：一个布尔值，如果为 True ，则返回完整的拟合信息。
# w ：一个数组，包含每个数据点的权重。
# cov ：一个布尔值，如果为 True ，则返回系数的协方差矩阵。
# 功能 ：
# np.polyfit 函数使用最小二乘法拟合一个 deg 阶多项式到数据点 (x, y) 。它返回多项式的系数，从最高次项到最低次项。
# 返回值 ：
# 函数返回一个数组，包含拟合多项式的系数。
# np.polyfit 是数据分析和科学计算中常用的函数，特别适用于趋势线分析和数据建模。通过拟合多项式，可以对数据进行预测或者提取数据的潜在规律。
# 使用 np.polyfit 函数对批量大小和内存使用情况进行一次多项式拟合。
p = np.polyfit(batch_sizes[: len(y)], y, deg=1) # first degree polynomial fit 一次多项式拟合。
# 计算最佳批量大小 b ，使得 CUDA 内存使用率达到 fraction 指定的百分比。
# 变量解释。
# f ：CUDA 设备的空闲内存（以吉字节为单位）。 fraction ：期望的 CUDA 内存使用率（以小数形式表示，例如 0.6 表示 60%）。 p ：一个包含多项式系数的元组，其中 p[0] 是多项式的斜率， p[1] 是 y 截距。
# 公式解释。
# f * fraction ：计算期望使用的 CUDA 内存量（以吉字节为单位）。 p[1] ：多项式 y 截距，表示在批量大小为 0 时的内存使用量。 p[0] ：多项式斜率，表示批量大小每增加一个单位时内存使用的增加量。 (f * fraction - p[1]) / p[0] ：计算达到期望内存使用率所需的批量大小的比率。 int(...) ：将计算结果转换为整数，因为批量大小必须是整数。
# 计算过程。
# 首先，计算期望使用的 CUDA 内存量 ： f * fraction 。 然后，从期望使用的内存量中减去多项式的 y 截距 ：f * fraction - p[1] 。 接着，将结果除以多项式的斜率 ：(f * fraction - p[1]) / p[0] 。 最后，将得到的比率转换为整数，得到最佳批量大小 b 。
# 这个公式基于多项式拟合的结果，通过解线性方程来找到在给定内存使用率下的最佳批量大小。这种方法允许在不超过 CUDA 内存限制的情况下，最大化模型训练的批量大小，从而提高训练效率。
b = int((f * fraction - p[1]) / p[0]) # y intercept (optimal batch size) y 截距（最佳批量大小）。
# 处理异常情况。
# 检查结果中是否包含 None。这行代码检查 results 列表中是否包含 None 值，这意味着在测试某个批量大小时出现了错误。
if None in results: # some sizes failed 有些尺寸不合格。
# 找到第一个失败的批量大小的索引。 如果发现 None ，则使用 index 方法找到第一个 None 值的索引 i 。
i = results.index(None) # first fail index 首次失败指数。
# 检查最佳批量大小是否大于或等于失败的批量大小。这行代码检查计算出的最佳批量大小 b 是否大于或等于失败的批量大小 batch_sizes[i] 。
if b >= batch_sizes[i]: # y intercept above failure point y 截距高于故障点。
# 选择一个安全的批量大小。 如果 b 大于或等于失败的批量大小，则选择一个安全的批量大小，即失败批量大小之前的最大的批量大小。这里使用 max(i - 1, 0) 确保索引值不会是负数。
b = batch_sizes[max(i - 1, 0)] # select prior safe point 选择先前的安全点。
# 如果最佳批量大小 b 在 1 到 1024 的范围之外，使用默认批量大小。
if b < 1 or b > 1024: # b outside of safe range b 超出安全范围。
b = batch_size
LOGGER.info(f"{prefix}WARNING ⚠️ CUDA anomaly detected, using default batch-size {batch_size}.") # {prefix}警告⚠️检测到 CUDA 异常，使用默认批量大小 {batch_size}。
# numpy.polyval(p, x)
# np.polyval 是 NumPy 库中的一个函数，用于计算多项式在特定值的函数值。
# 参数 :
# p ：一个数组，包含多项式的系数，从最高次项到常数项。
# x ：一个数值、数组或 poly1d 对象，表示要计算多项式值的点。
# 功能 :
# np.polyval 函数计算多项式 p 在点 x 处的值。如果 p 的长度为 N ，则函数返回的值按照以下公式计算 ：p[0]*x**(N-1) + p[1]*x**(N-2) + ... + p[N-2]*x + p[N-1]。
# 如果 x 是一个序列，则 p(x) 返回每个元素的多项式值。如果 x 是另一个多项式，则返回复合多项式 p(x(t)) 。
# 返回值 ：
# 函数返回一个 ndarray 或 poly1d 对象，包含在 x 处计算的多项式值。
# np.polyval 是一个方便的工具，用于在给定点评估多项式，常用于数据分析、科学计算和工程领域中的多项式函数计算。
# 计算实际内存使用率。 使用拟合的多项式计算实际的 CUDA 内存使用率。
fraction = (np.polyval(p, b) + r + a) / t # actual fraction predicted 实际分数 预测。
# 记录信息并返回最佳批量大小。
# 记录使用的最佳批量大小和实际的 CUDA 内存使用率。
LOGGER.info(f"{prefix}Using batch-size {b} for {d} {t * fraction:.2f}G/{t:.2f}G ({fraction * 100:.0f}%) ✅") # {prefix}使用批量大小 {b} 表示 {d} {t * 分数：.2f}G/{t：.2f}G ({分数 * 100：.0f}%)✅。
# 返回计算出的最佳批量大小。
return b
# 异常处理。
# 如果在尝试过程中出现任何异常，记录警告信息并返回默认批量大小。
except Exception as e:
LOGGER.warning(f"{prefix}WARNING ⚠️ error detected: {e}, using default batch-size {batch_size}.") # {prefix}警告 ⚠️ 检测到错误：{e}，使用默认批量大小 {batch_size}。
return batch_size
# 这段代码的目的是自动确定最佳批量大小，以便在不超过指定 CUDA 内存使用率的情况下最大化模型性能。通过这种方式，可以确保模型在有限的硬件资源下高效运行，同时避免内存溢出错误。
# 这个方法的目的是自动调整模型的批量大小，以便在不超过指定 CUDA 内存使用率的情况下最大化模型性能。通过这种方式，可以确保模型在有限的硬件资源下高效运行，同时避免内存溢出错误。