YOLOv8-ultralytics-8.2.103部分代码阅读笔记-autobatch.py

autobatch.py

ultralytics\utils\autobatch.py

目录

autobatch.py

1.所需的库和模块

2.def check_train_batch_size(model, imgsz=640, amp=True, batch=-1): 

3.def autobatch(model, imgsz=640, fraction=0.60, batch_size=DEFAULT_CFG.batch): 

1.所需的库和模块

# 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

2.def check_train_batch_size(model, imgsz=640, amp=True, batch=-1): 

# 这段代码定义了一个名为 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 内存使用率的情况下最大化模型性能。通过这种方式，可以确保模型在有限的硬件资源下高效运行，同时避免内存溢出错误。

3.def autobatch(model, imgsz=640, fraction=0.60, batch_size=DEFAULT_CFG.batch): 

# 这段代码定义了一个名为 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 内存使用率的情况下最大化模型性能。通过这种方式，可以确保模型在有限的硬件资源下高效运行，同时避免内存溢出错误。