YOLOv9-0.1部分代码阅读笔记-autobatch.py
autobatch.py
utils\autobatch.py
目录
autobatch.py
1.所需的库和模块
2.def check_train_batch_size(model, imgsz=640, amp=True):
3.def autobatch(model, imgsz=640, fraction=0.8, batch_size=16):
1.所需的库和模块
from copy import deepcopy
import numpy as np
import torch
from utils.general import LOGGER, colorstr
from utils.torch_utils import profile2.def check_train_batch_size(model, imgsz=640, amp=True):
# 这段代码定义了一个名为 check_train_batch_size 的函数，它用于计算给定模型在特定图像尺寸下的最优批量大小（batch size）。
# 定义了一个名为 check_train_batch_size 的函数，它接受三个参数。
# 1.model ：要检查的模型。
# 2.imgsz ：图像尺寸，默认为640。
# 3.amp ：一个布尔值，指示是否使用自动混合精度（Automatic Mixed Precision，AMP），默认为True。
def check_train_batch_size(model, imgsz=640, amp=True):
# Check YOLOv5 training batch size 检查 YOLOv5 训练批次大小。
# 使用 torch.cuda.amp.autocast 上下文管理器来启用或禁用自动混合精度（AMP）。如果 amp 参数为True，则启用AMP以提高训练速度和减少内存使用。
with torch.cuda.amp.autocast(amp):
# deepcopy(model) 创建模型的一个深拷贝，以避免修改原始模型。
# deepcopy(model).train() 将模型的副本设置为训练模式。
# autobatch(...) 函数被调用来计算最优批量大小。这个函数执行一系列操作来确定在给定图像尺寸和硬件条件下，模型可以处理的最大批量大小，而不会导致内存溢出或其他问题。
# imgsz 参数传递给 autobatch 函数，可能用于调整输入图像尺寸或作为计算批量大小的依据。
return autobatch(deepcopy(model).train(), imgsz) # compute optimal batch size
# 这个函数的目的是在不超出硬件资源限制的情况下，找到模型可以处理的最大批量大小，这对于优化训练过程和硬件资源利用非常重要。自动混合精度（AMP）的使用可以帮助在保持模型精度的同时减少内存消耗和加速训练。3.def autobatch(model, imgsz=640, fraction=0.8, batch_size=16):
# 这段代码定义了一个名为 autobatch 的函数，它用于自动计算模型在特定图像尺寸下的最优批量大小（batch size），以确保不超过 GPU 内存的限制。
# 定义了一个名为 autobatch 的函数，它接受四个参数。
# 1.model ：要检查的模型。
# 2.imgsz ：图像尺寸，默认为640。
# 3.fraction ：使用的 GPU 内存比例，默认为0.8。
# 4.batch_size ：默认的批量大小。
def autobatch(model, imgsz=640, fraction=0.8, batch_size=16):
# Automatically estimate best YOLOv5 batch size to use `fraction` of available CUDA memory 自动估计最佳 YOLOv5 批量大小以使用可用 CUDA 内存的“部分”。
# Usage:
# import torch
# from utils.autobatch import autobatch
# model = torch.hub.load('ultralytics/yolov5', 'yolov5s', autoshape=False)
# print(autobatch(model))
# Check device
# 使用 colorstr 函数来给日志前缀添加颜色，这里可能是为了在日志中突出显示。
# def colorstr(*input): -> 构建并返回最终的字符串。它首先通过列表推导式和 join 函数将所有颜色和样式的 ANSI 代码连接起来，然后加上要着色的字符串 string ，最后加上 colors['end'] 来重置样式，确保之后的输出不会受到颜色代码的影响。 -> return ''.join(colors[x] for x in args) + f'{string}' + colors['end']
prefix = colorstr('AutoBatch: ') # 自动批处理：
# 使用日志记录器 LOGGER 记录一条信息，表明正在计算最优批量大小。
LOGGER.info(f'{prefix}Computing optimal batch size for --imgsz {imgsz}') # {prefix}计算 --imgsz {imgsz} 的最佳批量大小。
# 获取模型参数所在的设备（GPU或CPU）。
device = next(model.parameters()).device # get model device
# 如果设备是CPU，则返回默认的批量大小。
if device.type == 'cpu':
LOGGER.info(f'{prefix}CUDA not detected, using default CPU batch-size {batch_size}') # {prefix}未检测到CUDA，使用默认CPU批次大小{batch_size}。
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，则建议将其设置为False，并返回默认的批量大小。
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}。
return batch_size
# Inspect CUDA memory
# 定义变量 gb 用于将字节转换为 GiB。
# 这行代码是一个位运算表达式，用于在Python中计算2的30次方，即1024的三次方。这个值等于1073741824，也就是1 GiB（Gibibyte），是计算机存储容量的一个单位。
# 1 是二进制数 00000000000000000000000000000001 。
# << 是位左移运算符，它将数字的二进制表示向左移动指定的位数。
# 30 表示将 1 的二进制表示向左移动30位。每次向左移动一位，数值都会翻倍，因此向左移动30位相当于乘以2的30次方，即： 1 × 2^30 = 1073741824 。
# 这个值通常用于将字节（Bytes）转换为吉字节（Gibibytes），因为1 GiB等于2的30次方字节。在处理文件大小和内存容量时，这种转换很有用。
gb = 1 << 30 # bytes to GiB (1024 ** 3)
# 获取设备名称，并转换为大写。
d = str(device).upper() # 'CUDA:0'
# torch.cuda.get_device_properties(device)
# torch.cuda.get_device_properties() 函数是 PyTorch 中用于获取指定 GPU 设备属性的函数。
# 参数 ：
# device (torch.device 或 int 或 str) ：要返回设备属性的设备。可以是设备 ID（整数），类似于 gpu:x 的设备名称，或者是 torch.device 对象。如果 device 为空，则默认为当前设备。默认值为 None。
# 返回值 ：
# 返回一个 _CudaDeviceProperties 对象，其中包含以下属性 ：
# name ：GPU 设备的名称（字符串）。
# total_memory ：GPU 总显存大小，以字节为单位（整数）。
# major ：CUDA 计算能力的主要版本号（整数）。
# minor ：CUDA 计算能力的次要版本号（整数）。
# multi_processor_count ：GPU 设备的多处理器数量（整数）。
# is_integrated ：GPU 是否是集成显卡，返回 True 或 False（布尔值）。
# is_multi_gpu_board ：GPU 是否是多 GPU 板卡，返回 True 或 False（布尔值）。
# 这个函数可以帮助你了解指定 GPU 设备的详细信息，例如显存大小、计算能力等，这些信息对于深度学习模型的训练和优化非常有用。
# 获取设备的属性。
properties = torch.cuda.get_device_properties(device) # device properties
# 计算总内存（GiB）。
t = properties.total_memory / gb # GiB total
# torch.cuda.memory_reserved(device=None)
# torch.cuda.memory_reserved() 是 PyTorch 提供的一个函数，用于查询当前 GPU 上由 PyTorch 预先保留但尚未实际分配的显存量。
# 参数 ：
# device (torch.device 或 str, 可选) ：指定要查询的 CUDA 设备。可以是设备索引（例如 'cuda:0' ），或者是 torch.device 对象。如果未指定，则默认为当前设备。
# 返回值 ：
# 返回一个整数，表示当前进程所分配的显存缓冲区总量，单位为字节。
# 功能描述 ：
# 这个函数返回由 PyTorch 预先保留但尚未实际分配的显存量。PyTorch 在启动时会预留一定数量的显存作为缓冲区，以便在需要时快速分配。这部分显存并不会直接被张量或模型参数占用，因此不计入 torch.cuda.memory_allocated() 的已分配显存。
# 使用场景 ：
# 监控和管理 GPU 显存使用情况，特别是在运行大型模型或处理大规模数据时，了解显存的预留情况对于优化性能和避免显存溢出非常重要。
# 注意事项 ：
# 返回的显存预留量包括 PyTorch 为当前进程预留的所有显存，但不包括其他进程可能占用的显存。
# torch.cuda.memory_reserved() 的值并不等同于 nvidia-smi 显示的值， nvidia-smi 显示的是 reserved_memory 和 PyTorch context 显存之和。
# 计算保留的内存（GiB）。
r = torch.cuda.memory_reserved(device) / gb # GiB reserved
# torch.cuda.memory_allocated(device=None)
# torch.cuda.memory_allocated() 是 PyTorch 提供的一个函数，用于查询当前 GPU 显存的使用情况。
# 参数 ：
# device (torch.device 或 int 或 str, 可选) ：指定要查询的 CUDA 设备。可以是设备索引（整数），类似于 cuda:x 的设备名称，或者是 torch.device 对象。如果未指定，则默认为当前设备。
# 返回值 ：
# 返回一个整数，表示已经由 PyTorch 分配但尚未释放的显存字节数。
# 功能描述 ：
# 这个函数返回已经由 PyTorch 分配但尚未释放的显存字节数。它只能反映出 PyTorch 所管理的显存使用情况，即已被张量、模型参数等数据占用的显存量。这个函数不包括由 PyTorch 预先保留但尚未实际分配的显存，这部分显存由 torch.cuda.memory_reserved() 函数返回。
# 使用场景 ：
# 监控和管理 GPU 显存使用情况，特别是在运行大型模型或处理大规模数据时，了解显存使用情况对于优化性能和避免显存溢出非常重要。
# 在调试时，检查显存分配是否符合预期，帮助定位显存泄漏问题。
# 注意事项 ：
# 返回的显存使用量仅限于当前 PyTorch 进程，不包括其他进程可能占用的显存。
# 该函数返回的值可能与实际显存使用情况有所差异，因为它不包括 PyTorch 预先保留的显存。
# 计算已分配的内存（GiB）。
a = torch.cuda.memory_allocated(device) / gb # GiB allocated
# 计算剩余的内存（GiB）。
f = t - (r + a) # GiB free
# 记录设备的内存信息。
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
# 定义一个列表，包含要测试的批量大小。
batch_sizes = [1, 2, 4, 8, 16]
# 尝试执行以下代码块。
try:
# 创建不同批量大小的空图像张量。
img = [torch.empty(b, 3, imgsz, imgsz) for b in batch_sizes]
# 调用 profile 函数来测试不同批量大小的内存使用情况。
# def profile(input, ops, n=10, device=None):
# -> 是一个性能分析工具，用于评估 PyTorch 模型或操作（ops）的性能。它计算给定输入（input）和操作（ops）的 参数数量 、 浮点运算次数（GFLOPs） 、 GPU显存使用量 、 前向传播时间 和 反向传播时间 。返回包含所有评估结果的列表 results 。
# -> return results
results = profile(img, model, n=3, device=device)
# 如果发生异常，记录警告信息。
except Exception as e:
LOGGER.warning(f'{prefix}{e}')
# Fit a solution
# 从结果中提取内存使用情况。
y = [x[2] for x in results if x] # memory [2]
# np.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 是数据分析和科学计算中常用的函数，特别适用于趋势线分析和数据建模。通过拟合多项式，可以对数据进行预测或者提取数据的潜在规律。
# 使用一阶多项式拟合批量大小和内存使用情况。
p = np.polyfit(batch_sizes[:len(y)], y, deg=1) # first degree polynomial fit
# 计算最优批量大小。
b = int((f * fraction - p[1]) / p[0]) # y intercept (optimal batch size)
# 这段代码是 autobatch 函数的一部分，用于处理在测试不同批量大小时可能遇到的问题，并确保最终选择的批量大小 b 是安全的。
# 检查 results 列表中是否包含 None 值，这意味着某些批量大小的测试失败了。
if None in results: # some sizes failed
# 如果有 None 值，找到第一个 None 值的索引 i ，这个索引对应于第一个失败的批量大小。
i = results.index(None) # first fail index
# 检查计算出的最优批量大小 b 是否大于或等于第一个失败的批量大小 batch_sizes[i] 。
if b >= batch_sizes[i]: # y intercept above failure point
# 如果最优批量大小 b 大于或等于失败的批量大小，那么选择前一个安全的批量大小作为最优批量大小。这里使用 max(i - 1, 0) 确保索引不会小于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 不在安全范围内，使用默认的批量大小 batch_size 作为最优批量大小。
b = batch_size
# 使用日志记录器 LOGGER 记录一条警告信息，提示检测到CUDA异常，建议重启环境并重试命令。
LOGGER.warning(f'{prefix}WARNING ⚠️ CUDA anomaly detected, recommend restart environment and retry command.') # {prefix}WARNING ⚠️ 检测到 CUDA 异常，建议重新启动环境并重试命令。
# 这段代码的目的是确保即使在某些批量大小测试失败或最优批量大小超出预期范围时，也能选择一个合理的批量大小。这样做可以避免因内存不足或其他问题导致的程序崩溃，并确保模型训练的稳定性。
# 计算实际使用的内存比例。
fraction = (np.polyval(p, b) + r + a) / t # actual fraction predicted
# 记录最终使用的批量大小和内存使用情况。
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
# 这个函数的目的是在不超出 GPU 内存限制的情况下，自动计算并返回最优的批量大小。它通过测试不同的批量大小并使用多项式拟合来预测内存使用情况，从而找到最优的批量大小。