autobatch.py
ultralytics\utils\autobatch.py
目录
autobatch.py
1.所需的库和模块
2.def check_train_batch_size(model, imgsz=640, amp=True, batch=-1, max_num_obj=1):
3.def autobatch(model, imgsz=640, fraction=0.60, batch_size=DEFAULT_CFG.batch, max_num_obj=1):
1.所需的库和模块
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
# 用于估算最佳 YOLO 批量大小的函数，以使用 PyTorch 中可用 CUDA 内存的一小部分。
"""Functions for estimating the best YOLO batch size to use a fraction of the available CUDA memory in PyTorch."""import os
from copy import deepcopyimport numpy as np
import torchfrom ultralytics.utils import DEFAULT_CFG, LOGGER, colorstr
from ultralytics.utils.torch_utils import autocast, profile2.def check_train_batch_size(model, imgsz=640, amp=True, batch=-1, max_num_obj=1):
# 这段代码定义了一个函数 check_train_batch_size ，用于检查和自动调整训练时的最佳批量大小。它通过调用 autobatch 函数来实现这一功能，并支持自动混合精度（AMP）训练。
# 定义了一个函数 check_train_batch_size ，接收以下参数 ：
# 1.model ：要检查的 PyTorch 模型。
# 2.imgsz ：输入图像的尺寸，默认为 640。
# 3.amp ：是否启用自动混合精度（AMP）训练，默认为 True 。
# 4.batch ：指定的批量大小，默认为 -1 ，表示自动选择最佳批量大小。
# 5.max_num_obj ：每个图像的预测数量，默认为 1，用于模拟训练时的预测显存占用。
def check_train_batch_size(model, imgsz=640, amp=True, batch=-1, max_num_obj=1):# 使用 autobatch() 函数计算最佳 YOLO 训练批次大小。# 注意：# 如果 0.0 < batch < 1.0，则将其用作要使用的 GPU 内存分数。# 否则，使用默认分数 0.6。"""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.max_num_obj (int, optional): The maximum number of objects from dataset.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."""# 使用 autocast 上下文管理器，启用或禁用自动混合精度（AMP）训练。如果 amp 为 True ，则启用 AMP，否则禁用。# def autocast(enabled: bool, device: str = "cuda"):# -> 用于根据 PyTorch 的版本动态选择混合精度训练的上下文管理器。如果当前 PyTorch 版本为 1.13 或更高版本，返回 torch.amp.autocast 上下文管理器。返回 torch.cuda.amp.autocast ，并传递 enabled 参数，用于控制是否启用混合精度训练。# -> return torch.amp.autocast(device, enabled=enabled) / return torch.cuda.amp.autocast(enabled)with autocast(enabled=amp):# 调用 autobatch 函数，传入以下参数。return autobatch(# deepcopy(model).train() ：创建模型的一个深拷贝，并将其设置为训练模式。这一步是为了避免对原始模型进行修改。# imgsz ：输入图像的尺寸。# fraction=batch if 0.0 < batch < 1.0 else 0.6 ：指定批量大小的分数。如果 batch 在 (0.0, 1.0) 范围内，则直接使用 batch ；否则，默认使用 0.6 。# max_num_obj ：每个图像的预测数量，用于模拟训练时的预测显存占用。# autobatch 函数的作用是自动调整批量大小，以确保在训练过程中不会因显存不足而失败。它会根据显存占用情况动态调整批量大小，找到一个合适的值。deepcopy(model).train(), imgsz, fraction=batch if 0.0 < batch < 1.0 else 0.6, max_num_obj=max_num_obj)
# check_train_batch_size 函数通过调用 autobatch 函数，自动调整训练时的最佳批量大小。它支持自动混合精度（AMP）训练，并允许用户指定输入图像尺寸、批量大小的分数以及每个图像的预测数量。通过这种方式，用户可以更灵活地调整训练参数，确保训练过程的顺利进行，同时充分利用显存资源。3.def autobatch(model, imgsz=640, fraction=0.60, batch_size=DEFAULT_CFG.batch, max_num_obj=1):
# 这段代码定义了一个名为 autobatch 的函数，用于自动调整训练时的最佳批量大小，以充分利用 GPU 显存，同时避免显存溢出。它通过分析不同批量大小下的显存占用情况，动态选择一个合适的批量大小。
# 定义了一个函数 autobatch ，接收以下参数 ：
# 1.model ：要检查的 PyTorch 模型。
# 2.imgsz ：输入图像的尺寸，默认为 640。
# 3.fraction ：目标显存利用率，默认为 0.60，表示目标显存占用率为 60%。
# 4.batch_size ：默认批量大小，用于在自动调整失败时回退。
# 5.max_num_obj ：每个图像的预测数量，默认为 1，用于模拟训练时的预测显存占用。
def autobatch(model, imgsz=640, fraction=0.60, batch_size=DEFAULT_CFG.batch, max_num_obj=1):# 自动估计最佳 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.max_num_obj (int, optional): The maximum number of objects from dataset.Returns:(int): The optimal batch size."""# 这段代码是 autobatch 函数的起始部分，主要用于检查设备类型和配置，确保函数适用于 CUDA 设备，并在不满足条件时直接返回默认批量大小。# Check device# 定义一个前缀字符串 prefix ，用于日志输出。 colorstr 是一个函数，用于为字符串添加颜色，便于在终端中突出显示。这里将 "AutoBatch: " 作为前缀，用于标识日志信息的来源。prefix = colorstr("AutoBatch: ")# 使用 LOGGER.info 打印一条日志信息，说明正在计算目标显存利用率下的最佳批量大小。日志中包含输入图像尺寸 imgsz 和目标显存利用率（以百分比表示）。LOGGER.info(f"{prefix}Computing optimal batch size for imgsz={imgsz} at {fraction * 100}% CUDA memory utilization.") # {prefix}在 CUDA 内存利用率为 {fraction * 100}% 时，计算 imgsz={imgsz} 的最佳批量大小。# 获取模型的设备信息。通过访问模型的第一个参数的 .device 属性，确定模型当前运行的设备（如 CPU、GPU 或 MPS）。device = next(model.parameters()).device # get model device# 检查设备类型是否为 CPU 或 MPS。if device.type in {"cpu", "mps"}:# 如果模型运行在这些设备上，而不是 CUDA 设备上，则打印警告信息。LOGGER.info(f"{prefix} ⚠️ intended for CUDA devices, using default batch-size {batch_size}") # {prefix} ⚠️适用于 CUDA 设备，使用默认批量大小 {batch_size}。# 并直接返回默认批量大小 batch_size 。这是因为 autobatch 函数主要针对 CUDA 设备设计，用于动态调整 GPU 上的批量大小。return batch_size# torch.backends.cudnn.benchmark# cudnn.benchmark 是 PyTorch 中的一个设置，用于控制 NVIDIA 的 cuDNN 库是否在程序运行时自动为每个卷积层选择最优的算法。这个设置可以影响程序的性能，尤其是在深度学习模型中使用卷积层时。# 定义和用法 ：# torch.backends.cudnn.benchmark ：这是一个布尔值设置，可以设置为 True 或 False 。# True ：开启 cuDNN 的基准测试模式。在这个模式下，cuDNN 会在程序开始运行时为每个卷积层自动选择最优的算法。这可能会在程序启动时增加一些额外的时间开销，因为 cuDNN 需要对不同的算法进行基准测试，但一旦选择了最优算法，后续的卷积操作将会更快。# False ：关闭基准测试模式。cuDNN 将使用默认的卷积算法，这可能不是最优的选择，但适用于模型输入尺寸在运行过程中会改变的情况。# 适用场景 ：# 固定输入尺寸 ：如果你的模型输入尺寸（例如，图像尺寸和批处理大小）是固定的，设置 torch.backends.cudnn.benchmark = True 可以提高运行效率，因为 cuDNN 可以预先选择最优算法。# 变化输入尺寸 ：如果输入尺寸可能发生变化，开启 benchmark 可能导致性能下降，因为每次输入尺寸改变时，cuDNN 都可能重新搜索算法。# 注意事项 ：# 性能影响 ：开启 cudnn.benchmark 可能会在程序启动时增加一些额外的时间开销，但可以提高后续卷积操作的速度。# 结果可重复性 ：开启 cudnn.benchmark 可能会导致结果的轻微变化，因为 cuDNN 可能会选择不同的算法。如果需要确保结果的完全可重复性，可能需要关闭 cudnn.benchmark 并设置 torch.backends.cudnn.deterministic = True 。# 总的来说， cudnn.benchmark 是一个有用的设置，可以帮助优化深度学习模型的性能，但需要根据具体的应用场景和需求来决定是否开启。# 检查 torch.backends.cudnn.benchmark 是否为 True 。if torch.backends.cudnn.benchmark:# 如果为 True ，则打印警告信息。LOGGER.info(f"{prefix} ⚠️ Requires torch.backends.cudnn.benchmark=False, using default batch-size {batch_size}") # {prefix} ⚠️ 需要 torch.backends.cudnn.benchmark=False，使用默认批量大小 {batch_size}。# 并直接返回默认批量大小 batch_size 。这是因为启用 cudnn.benchmark 时，PyTorch 会尝试优化卷积操作的性能，但可能会导致显存占用不稳定，从而影响批量大小的自动调整。return batch_size# 这段代码的主要作用是检查当前设备是否适合运行 autobatch 函数。它确保模型运行在 CUDA 设备上，并且 torch.backends.cudnn.benchmark 未启用。如果不满足这些条件，则直接返回默认批量大小，避免在不支持的设备上执行批量大小的自动调整。这种检查机制可以提高函数的健壮性，确保其在合适的环境下运行。# 这段代码的作用是检查当前 CUDA 设备的显存状态，包括总显存、已保留显存、已分配显存和空闲显存。它通过 PyTorch 提供的显存管理接口获取这些信息，并打印详细的显存状态日志。# Inspect CUDA memory# 定义一个变量 gb ，表示字节到 GiB 的转换因子。 1 << 30 等价于 1024^3 ，即 1 GiB 的字节数。gb = 1 << 30 # bytes to GiB (1024 ** 3)# 获取当前可见的 CUDA 设备编号。通过 os.getenv('CUDA_VISIBLE_DEVICES', '0') 获取环境变量 CUDA_VISIBLE_DEVICES 的值，如果没有设置，则默认为 '0' 。然后取其第一个字符（假设只使用一个 GPU），并格式化为字符串 'CUDA:0' 。d = f"CUDA:{os.getenv('CUDA_VISIBLE_DEVICES', '0').strip()[0]}" # 'CUDA:0'# 调用 torch.cuda.get_device_properties(device) 获取当前设备的属性信息，包括设备名称、总显存等。properties = torch.cuda.get_device_properties(device) # device properties# 获取设备的总显存大小（以字节为单位），并将其转换为 GiB（通过除以 gb ）。t = properties.total_memory / gb # GiB total# 获取 当前设备的已保留显存大小（以字节为单位） ，并将其转换为 GiB。已保留显存是指已分配给 PyTorch 的显存，但尚未使用的部分。r = torch.cuda.memory_reserved(device) / gb # GiB reserved# 获取 当前设备的已分配显存大小（以字节为单位） ，并将其转换为 GiB。已分配显存是指当前正在使用的显存部分。a = torch.cuda.memory_allocated(device) / gb # GiB allocated# 计算 当前设备的空闲显存大小 。空闲显存等于总显存减去已保留显存和已分配显存。f = t - (r + a) # GiB free# 使用 LOGGER.info 打印一条日志信息，显示当前设备的显存状态。日志中包含 设备编号 、 设备名称 、 总显存 、 已保留显存 、 已分配显存 和 空闲显存 。LOGGER.info(f"{prefix}{d} ({properties.name}) {t:.2f}G total, {r:.2f}G reserved, {a:.2f}G allocated, {f:.2f}G free")# 这段代码通过调用 PyTorch 的显存管理接口，检查当前 CUDA 设备的显存状态，并打印详细的显存信息。这有助于开发者了解当前设备的显存使用情况，从而更好地进行显存管理和性能优化。# 这段代码是 autobatch 函数中用于测试不同批量大小下的显存占用情况的部分。它通过生成一系列不同批量大小的输入张量，并调用 profile 函数来测量每个批量大小下的显存占用和性能指标。# Profile batch sizes# 根据 GPU 的总显存大小 t （单位为 GiB），选择要测试的批量大小列表。# 如果总显存小于 16 GiB，则测试批量大小 [1, 2, 4, 8, 16] 。# 如果总显存大于或等于 16 GiB，则测试批量大小 [1, 2, 4, 8, 16, 32, 64] 。# 这种选择策略是为了确保在显存较小的设备上不会因测试过大的批量大小而导致显存溢出。batch_sizes = [1, 2, 4, 8, 16] if t < 16 else [1, 2, 4, 8, 16, 32, 64]# 使用列表推导式生成一个 输入张量列表 img ，每个张量的批量大小分别为 batch_sizes 中的值。# torch.empty(b, 3, imgsz, imgsz) 创建一个形状为 (b, 3, imgsz, imgsz) 的空张量，表示批量大小为 b 的输入图像数据。# imgsz 是输入图像的尺寸，默认为 640。# 这些张量将用于后续的性能分析。try:img = [torch.empty(b, 3, imgsz, imgsz) for b in batch_sizes]# 调用 profile 函数，对每个批量大小的输入张量进行性能分析。# img ：输入张量列表，每个张量对应一个批量大小。# model ：要分析的 PyTorch 模型。# n=1 ：表示每个批量大小只测试一次。 profile 函数通常会多次测试以获取更稳定的性能指标，但在这里只需要一次测试即可。# device=device ：指定设备，确保测试在正确的设备上进行。# max_num_obj=max_num_obj ：指定每个图像的预测数量，用于模拟训练时的预测显存占用。# results ： profile 函数返回的性能分析结果列表，每个元素包含对应批量大小的显存占用、前向传播时间和反向传播时间等信息。# def profile(input, ops, n=10, device=None, max_num_obj=0): -> 用于对 PyTorch 模型或操作进行性能分析，包括速度、显存占用和 FLOPs（浮点运算次数）。返回性能分析结果列表 results 。包括 参数数量 、 FLOPs 、 显存占用 、 前向传播时间 、 反向传播时间 、 输入形状 和 输出形状 。 -> return resultsresults = profile(img, model, n=1, device=device, max_num_obj=max_num_obj)# 功能说明 ：这段代码的核心目的是通过测试不同批量大小下的显存占用情况，为后续的批量大小自动调整提供数据支持。 profile 函数会记录每个批量大小下的显存占用、前向传播时间和反向传播时间等关键性能指标，这些数据将用于拟合显存占用与批量大小之间的关系，从而计算出最佳批量大小。# 这段代码通过生成一系列不同批量大小的输入张量，并调用 profile 函数测量每个批量大小下的显存占用和性能指标。它根据 GPU 的总显存大小选择合适的批量大小范围，确保测试过程不会因显存不足而失败。这些性能分析结果将为后续的批量大小自动调整提供重要的数据支持。# 这段代码是 autobatch 函数的核心部分，用于根据性能分析结果拟合显存占用与批量大小之间的关系，并计算出最佳批量大小。# Fit a solution# 构建一个列表 xy ，其中每个元素是一个二元组 [batch_size, memory_usage] ，表示 每个批量大小 及其 对应的显存占用 。xy = [[x, y[2]]for i, (x, y) in enumerate(zip(batch_sizes, results))# 过滤条件。# y 是有效的结果（非 None ）。if y # valid result# y[2] 是数值类型（显存占用）。and isinstance(y[2], (int, float)) # is numeric# 显存占用在合理范围内（大于 0 且小于 GPU 总显存 t ）。and 0 < y[2] < t # between 0 and GPU limit# 显存占用是递增的（或为第一个数据点）。# 这行代码是 autobatch 函数中用于过滤有效显存占用数据的一部分条件。它的作用是确保显存占用是递增的，或者当前批量大小是第一个测试点。# i == 0 ：表示当前批量大小是第一个测试点。对于第一个测试点，没有前一个数据点可以比较，因此直接接受该数据点。# not results[i - 1] ：表示前一个批量大小的测试结果是 None ，即前一个测试失败了。如果前一个测试失败，当前测试点仍然有效，因为失败点之前的最后一个有效数据点是有意义的。# y[2] > results[i - 1][2] ：表示当前批量大小的显存占用 y[2] 大于前一个批量大小的显存占用 results[i - 1][2] 。这确保了显存占用是递增的，符合显存占用随批量大小增加而增加的预期。# 这个条件的作用是过滤掉那些不符合显存占用递增规律的数据点。在实际测试中，显存占用通常会随着批量大小的增加而增加。如果某个批量大小的显存占用没有增加，或者比前一个批量大小的显存占用还小，那么这个数据点可能是异常的，应该被过滤掉。# 这行代码通过确保显存占用是递增的，或者当前批量大小是第一个测试点，过滤掉异常数据点，从而提高拟合结果的准确性和可靠性。这对于后续计算最佳批量大小至关重要。and (i == 0 or not results[i - 1] or y[2] > results[i - 1][2]) # first item or increasing memory]# 如果 xy 不为空，则将 xy 分解为两个列表 fit_x 和 fit_y ，分别表示 批量大小 和 显存占用 。 如果 xy 为空，则 fit_x 和 fit_y 均为空列表。fit_x, fit_y = zip(*xy) if xy else ([], [])# 使用 np.polyfit 进行一阶多项式拟合，拟合在对数空间中进行。拟合的目的是找到显存占用与批量大小之间的关系。 np.log(fit_x) 和 np.log(fit_y) 将数据转换为对数空间， deg=1 表示进行一阶多项式拟合。p = np.polyfit(np.log(fit_x), np.log(fit_y), deg=1) # first-degree polynomial fit in log space# 根据拟合结果计算 目标显存利用率下的最佳批量大小 b 。计算公式为 ：# b = exp((log(f×fraction)-p[1]))/p[0])# 其中 ：# f 是空闲显存（单位为 GiB）。# fraction 是目标显存利用率。# p[0] 和 p[1] 是拟合多项式的系数。b = int(round(np.exp((np.log(f * fraction) - p[1]) / p[0]))) # y intercept (optimal batch size)# 如果某些批量大小的测试失败（ results 中包含 None ），则找到第一个失败的批量大小索引 i 。if None in results: # some sizes failedi = results.index(None) # first fail index# 如果计算的批量大小 b 大于或等于失败点的批量大小，则选择失败点之前的最后一个安全批量大小。if b >= batch_sizes[i]: # y intercept above failure pointb = batch_sizes[max(i - 1, 0)] # select prior safe point# 如果计算的批量大小 b 不在安全范围内（小于 1 或大于 1024）。if b < 1 or b > 1024: # b outside of safe range# 则打印警告信息。LOGGER.info(f"{prefix}WARNING ⚠️ batch={b} outside safe range, using default batch-size {batch_size}.") # {prefix}警告 ⚠️ batch={b} 超出安全范围，使用默认批量大小 {batch_size}。# 并使用默认批量大小 batch_size 。b = batch_size# 这段代码通过拟合显存占用与批量大小之间的关系，计算出目标显存利用率下的最佳批量大小。它考虑了显存占用的递增性和测试失败的情况，并在计算结果不合理时回退到默认批量大小。这种方法可以动态调整批量大小，充分利用 GPU 显存，同时避免显存溢出。# 这段代码是 autobatch 函数的最后部分，用于计算最终的显存利用率、打印日志信息，并返回最佳批量大小。同时，它还处理了可能发生的异常，并在函数退出时清空 CUDA 缓存。# 计算 实际显存利用率 fraction 。# np.polyval(p, np.log(b)) ：根据拟合的多项式系数 p 和对数批量大小 np.log(b) ，计算预测的显存占用（对数空间）。# np.exp(...) ：将预测的显存占用从对数空间转换回线性空间。# r + a ：加上已保留显存 r 和已分配显存 a ，得到总的显存占用。# t ：GPU 的总显存。# fraction ：预测的显存利用率，表示为总显存的比例。fraction = (np.exp(np.polyval(p, np.log(b))) + r + a) / t # predicted fraction# 使用 LOGGER.info 打印一条日志信息，显示最终选择的批量大小及其显存占用情况。# b ：最佳批量大小。# d ：设备编号（如 CUDA:0 ）。# t * fraction ：预测的显存占用量（单位为 GiB）。# t ：GPU 的总显存（单位为 GiB）。# fraction * 100 ：预测的显存利用率（百分比）。LOGGER.info(f"{prefix}Using batch-size {b} for {d} {t * fraction:.2f}G/{t:.2f}G ({fraction * 100:.0f}%) ✅") # {prefix}使用批量大小 {b} 表示 {d} {t * fraction：.2f}G/{t：.2f}G ({fraction * 100：.0f}%)✅。# 返回计算的最佳批量大小 b 。return b# 如果在自动调整过程中发生异常。except Exception as e:# 则捕获异常并打印警告信息，说明错误原因。LOGGER.warning(f"{prefix}WARNING ⚠️ error detected: {e}, using default batch-size {batch_size}.") # {prefix}警告 ⚠️ 检测到错误：{e}，使用默认批量大小 {batch_size}。# 并返回默认批量大小 batch_size 。return batch_size# 在函数退出时，清空 CUDA 缓存，释放未使用的显存。这一步是为了确保在函数执行完毕后，显存被正确释放，避免显存泄漏。finally:torch.cuda.empty_cache()# 这段代码通过计算实际显存利用率并打印详细的日志信息，向用户展示最终选择的批量大小及其显存占用情况。它还处理了可能发生的异常，并在函数退出时清空 CUDA 缓存，确保资源被正确管理。这种设计提高了函数的健壮性和用户体验。
# autobatch 函数是一个用于自动调整训练时最佳批量大小的工具，旨在充分利用 GPU 显存资源，同时避免显存溢出。它通过分析不同批量大小下的显存占用情况，结合目标显存利用率，动态计算出一个合适的批量大小。该函数支持自动混合精度（AMP）训练，并在计算过程中考虑了显存利用率目标、显存占用的递增性以及测试失败的情况。最终，它返回一个优化后的批量大小，确保训练过程的高效性和稳定性。
