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）訓練，并在計算過程中考慮了顯存利用率目標、顯存占用的遞增性以及測試失敗的情況。最終，它返回一個優化后的批量大小，確保訓練過程的高效性和穩定性。
![Python3+Request+Pytest+Allure+Jenkins 接口自動化測試[手動寫的和AI寫的對比]](http://pic.xiahunao.cn/Python3+Request+Pytest+Allure+Jenkins 接口自動化測試[手動寫的和AI寫的對比])
