torch_utils.py
ultralytics\utils\torch_utils.py
目录
torch_utils.py
1.所需的库和模块
2.def torch_distributed_zero_first(local_rank: int):
3.def smart_inference_mode():
4.def get_cpu_info():
5.def select_device(device="", batch=0, newline=False, verbose=True):
6.def time_sync():
7.def fuse_conv_and_bn(conv, bn):
8.def fuse_deconv_and_bn(deconv, bn):
9.def model_info(model, detailed=False, verbose=True, imgsz=640):
10.def get_num_params(model):
11.def get_num_gradients(model):
12.def model_info_for_loggers(trainer):
13.def get_flops(model, imgsz=640):
14.def get_flops_with_torch_profiler(model, imgsz=640):
15.def initialize_weights(model):
16.def scale_img(img, ratio=1.0, same_shape=False, gs=32):
17.def make_divisible(x, divisor):
18.def copy_attr(a, b, include=(), exclude=()):
19.def get_latest_opset():
20.def intersect_dicts(da, db, exclude=()):
21.def is_parallel(model):
22.def de_parallel(model):
23.def one_cycle(y1=0.0, y2=1.0, steps=100):
24.def init_seeds(seed=0, deterministic=False):
25.class ModelEMA:
26.def strip_optimizer(f: Union[str, Path] = "best.pt", s: str = "") -> None:
27.def profile(input, ops, n=10, device=None):
28.class EarlyStopping:
1.所需的库和模块
# Ultralytics YOLO 🚀, AGPL-3.0 licenseimport math
import os
import random
import time
from contextlib import contextmanager
from copy import deepcopy
from pathlib import Path
from typing import Unionimport numpy as np
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.nn.functional as F
import torchvisionfrom ultralytics.utils import DEFAULT_CFG_DICT, DEFAULT_CFG_KEYS, LOGGER, __version__
from ultralytics.utils.checks import PYTHON_VERSION, check_versiontry:import thop
except ImportError:thop = None# Version checks (all default to version>=min_version)
# 这段代码使用 check_version 函数来检查 PyTorch 和 Torchvision 库的版本是否满足特定的要求。每个变量 TORCH_1_9 , TORCH_1_13 , TORCH_2_0 , TORCHVISION_0_10 , TORCHVISION_0_11 , 和 TORCHVISION_0_13 存储了一个布尔值,表示当前环境中相应版本的 PyTorch 或 Torchvision 是否可用。
# def check_version(current: str = "0.0.0", required: str = "0.0.0", name: str = "version", hard: bool = False, verbose: bool = False, msg: str = "",) -> bool:
# -> 用于检查当前安装的软件包版本是否满足特定的要求。返回 result ,表示版本检查是否通过。返回 False ,表示版本检查未通过。
# -> return result / return False
# 调用 check_version 函数检查当前安装的 PyTorch 版本是否至少为 1.9.0 版本。结果存储在 TORCH_1_9 变量中。
TORCH_1_9 = check_version(torch.__version__, "1.9.0")
# 调用 check_version 函数检查当前安装的 PyTorch 版本是否至少为 1.13.0 版本。结果存储在 TORCH_1_13 变量中。
TORCH_1_13 = check_version(torch.__version__, "1.13.0")
# 调用 check_version 函数检查当前安装的 PyTorch 版本是否至少为 2.0.0 版本。结果存储在 TORCH_2_0 变量中。
TORCH_2_0 = check_version(torch.__version__, "2.0.0")
# 调用 check_version 函数检查当前安装的 Torchvision 版本是否至少为 0.10.0 版本。结果存储在 TORCHVISION_0_10 变量中。
TORCHVISION_0_10 = check_version(torchvision.__version__, "0.10.0")
# 调用 check_version 函数检查当前安装的 Torchvision 版本是否至少为 0.11.0 版本。结果存储在 TORCHVISION_0_11 变量中。
TORCHVISION_0_11 = check_version(torchvision.__version__, "0.11.0")
# 调用 check_version 函数检查当前安装的 Torchvision 版本是否至少为 0.13.0 版本。结果存储在 TORCHVISION_0_13 变量中。
TORCHVISION_0_13 = check_version(torchvision.__version__, "0.13.0")
# 这些变量可以在代码中用于条件逻辑,以确保只有在满足特定版本要求时才执行某些操作。例如,如果代码依赖于 PyTorch 1.9.0 版本中引入的特性,那么可以检查 TORCH_1_9 变量的值来决定是否执行相关代码或者提供错误消息。
2.def torch_distributed_zero_first(local_rank: int):
# 这段代码定义了一个名为 torch_distributed_zero_first 的上下文管理器,用于在分布式训练环境中确保只有本地排名为0的进程首先执行某些操作,其他进程等待。这在分布式训练中非常有用,特别是在需要确保只有主进程(通常是排名为0的进程)进行某些操作(如保存模型或日志)时。# @contextlib.contextmanager
# contextlib.contextmanager 是 Python 标准库 contextlib 模块中的一个装饰器,它用于将一个生成器函数转换为上下文管理器。上下文管理器是一种特殊的对象,它允许你定义在代码块执行前后需要运行的代码,通常用于管理资源,如文件操作、锁的获取和释放等。
# 装饰器定义 :
# @contextlib.contextmanager
# def some_context_manager():
# # 在代码块执行前需要执行的代码
# yield # 这表示上下文管理器的入口点
# # 在代码块执行后需要执行的代码
# 函数体 :
# yield 语句之前的代码块在进入 with 语句时执行,这通常用于设置或初始化资源。
# yield 语句之后的代码块在退出 with 语句时执行,这通常用于清理或释放资源。
# contextlib.contextmanager 装饰器提供了一种简洁的方式来创建上下文管理器,而不需要定义一个类并实现 __enter__ 和 __exit__ 方法。这种方式特别适合于简单的资源管理场景,可以减少模板代码,使代码更加清晰和易于维护。它特别适用于需要执行特定进入和退出操作的场景,如文件操作、数据库连接、网络请求等。@contextmanager
# 使用 @contextmanager 装饰器定义了一个上下文管理器 torch_distributed_zero_first ,它接受一个参数。
# 1.local_rank :当前进程的本地排名。
def torch_distributed_zero_first(local_rank: int):# 装饰器使分布式训练中的所有进程等待每个 local_master 执行某项操作。"""Decorator to make all processes in distributed training wait for each local_master to do something."""# torch.distributed.is_available() -> bool# torch.distributed.is_available() 是PyTorch库中的一个函数,用于检查分布式训练功能是否在当前环境中可用。这个函数返回一个布尔值,表示分布式训练功能是否可以使用。# 返回值 :# bool ,表示分布式训练功能是否可用。# 用途 :# 检查分布式训练支持 :# 在使用分布式训练功能之前,通常需要检查当前环境是否支持分布式训练。 torch.distributed.is_available() 可以用于此目的。# 如果返回 True ,表示当前环境支持分布式训练,可以继续使用分布式训练的相关功能。# 如果返回 False ,表示当前环境不支持分布式训练,可能需要检查环境配置或安装相关的依赖。# 条件执行 :# 在编写分布式训练代码时,可以使用 torch.distributed.is_available() 来条件执行某些代码块,确保只有在支持分布式训练的环境中才执行相关操作。# torch.distributed.is_available() 是一个简单的检查函数,用于确定当前环境是否支持分布式训练。在使用分布式训练功能之前,建议先调用这个函数,以确保环境配置正确,避免在不支持分布式训练的环境中运行相关代码。# torch.distributed.is_initialized() -> bool# torch.distributed.is_initialized() 是PyTorch中的一个函数,用于检查分布式训练环境是否已经初始化。这个函数返回一个布尔值,表示分布式训练环境是否已经准备好。# 返回值 :# bool ,表示分布式训练环境是否已经初始化。# 用途 :# 检查初始化状态 :# 在使用分布式训练功能之前,通常需要检查分布式训练环境是否已经初始化。 torch.distributed.is_initialized() 可以用于此目的。# 如果返回 True ,表示分布式训练环境已经初始化,可以继续使用分布式训练的相关功能。# 如果返回 False ,表示分布式训练环境尚未初始化,需要调用 torch.distributed.init_process_group() 来初始化。# 条件执行 :# 在编写分布式训练代码时,可以使用 torch.distributed.is_initialized() 来条件执行某些代码块,确保只有在分布式训练环境已经初始化的情况下才执行相关操作。# torch.distributed.is_initialized() 是一个简单的检查函数,用于确定分布式训练环境是否已经初始化。在使用分布式训练功能之前,建议先调用这个函数,以确保环境配置正确,避免在未初始化的环境中运行相关代码。通过条件执行,可以确保只有在分布式训练环境已经初始化的情况下才执行特定的操作,提高代码的健壮性和可靠性。# 检查分布式训练环境是否已初始化。 torch.distributed.is_available() 检查分布式训练是否可用, torch.distributed.is_initialized() 检查分布式训练是否已初始化。如果两者都为 True ,则 initialized 为 True 。initialized = torch.distributed.is_available() and torch.distributed.is_initialized()# 如果分布式训练已初始化且当前进程的本地排名不是-1或0,则调用 dist.barrier(device_ids=[local_rank]) 。这会阻塞当前进程,直到所有指定的进程到达这个屏障。这里使用 device_ids=[local_rank] 确保只有当前进程参与屏障同步。if initialized and local_rank not in (-1, 0):# dist.barrier(device_ids=[device_id1, device_id2, ...])# 在 PyTorch 的分布式训练中, dist.barrier() 函数用于同步所有进程,确保它们在继续执行之前都达到了相同的执行点。当涉及到多个设备(例如,多个 GPU)时, dist.barrier() 可以接受一个 device_ids 参数,用于指定哪些设备参与同步。# 参数 :# device_ids :(可选)一个包含设备 ID 的列表。如果提供, dist.barrier() 将只同步指定的设备。如果不提供或为 None ,则同步所有设备。# 工作原理 :# 当一个进程中的某个设备调用 dist.barrier(device_ids) 后,它会暂停执行,直到所有进程中对应 device_ids 指定的设备都调用了该函数。# 一旦所有指定的设备都到达了 barrier 点, dist.barrier() 函数会返回,然后所有设备继续执行后续的操作。# 特点 :# dist.barrier() 函数通常用于确保在分布式训练中所有进程都完成了某个特定的操作后再进行后续的操作,从而保证训练的一致性。# 它可以用来同步数据加载、模型参数更新等操作,确保所有进程在执行下一步之前都已经完成了当前步骤。# 在使用 dist.barrier() 时,需要先初始化进程组,并在结束时销毁进程组。# dist.barrier(device_ids) 函数是 PyTorch 分布式训练中实现设备间同步的重要工具,它确保了在执行下一步之前,所有指定的设备都已经完成了当前步骤,从而保证了训练的一致性和稳定性。dist.barrier(device_ids=[local_rank])# yield 语句是上下文管理器的关键部分,它表示上下文管理器的“主体”部分。在这个点上,上下文管理器会暂停执行,直到上下文块中的代码执行完毕。这允许在 with 语句中执行用户代码。yield# 如果分布式训练已初始化且当前进程的本地排名为0,则调用 dist.barrier(device_ids=[0]) 。这会阻塞排名为0的进程,直到所有其他进程到达这个屏障。这里使用 device_ids=[0] 确保只有排名为0的进程参与屏障同步。if initialized and local_rank == 0:dist.barrier(device_ids=[0])
# torch_distributed_zero_first 上下文管理器确保在分布式训练环境中,只有本地排名为0的进程首先执行某些操作,其他进程等待。这在需要确保只有主进程进行特定操作(如保存模型或日志)时非常有用。通过使用 dist.barrier ,可以同步所有进程,确保操作的顺序性和一致性。
3.def smart_inference_mode():
# 这段代码定义了一个名为 smart_inference_mode 的函数,其作用是根据PyTorch的版本号来选择性地应用 torch.inference_mode() 装饰器或 torch.no_grad() 装饰器,以优化模型的推理模式。
# 这行代码定义了一个名为 smart_inference_mode 的函数,它不接受任何参数。
def smart_inference_mode():# 如果 torch>=1.9.0,则应用 torch.inference_mode() 装饰器,否则应用 torch.no_grad() 装饰器。"""Applies torch.inference_mode() decorator if torch>=1.9.0 else torch.no_grad() decorator."""# 在 smart_inference_mode 函数内部,定义了一个嵌套函数 decorate ,它接受一个参数。# 1.fn :这个参数是一个函数,即要被装饰的函数。def decorate(fn):# 根据 torch 版本为推理模式应用适当的 torch 装饰器。"""Applies appropriate torch decorator for inference mode based on torch version."""# torch.is_inference_mode_enabled()# torch.is_inference_mode_enabled() 是 PyTorch 提供的一个函数,用于检查当前是否启用了推理模式(inference mode)。当在推理模式下运行代码时,PyTorch 会禁用梯度计算和某些其他功能,以提高性能。# 返回值 :# 返回一个布尔值( bool ),如果当前启用了推理模式,则返回 True ;否则返回 False 。# 相关背景 :# 从 PyTorch 1.9 版本开始,引入了 torch.inference_mode 上下文管理器,它类似于 torch.no_grad ,但在推理模式下运行的代码可以通过禁用视图跟踪和版本计数来获得更好的性能。# torch.inference_mode 可以作为上下文管理器使用,也可以作为装饰器使用。# 上下文管理器示例:# with torch.inference_mode():# 在这个示例中,使用 torch.inference_mode 上下文管理器来禁用梯度计算,然后执行乘法操作。在推理模式下, y 的 requires_grad 属性将被设置为 False 。# 装饰器示例:# @torch.inference_mode()# 在这个示例中,使用 @torch.inference_mode() 装饰器来装饰函数 func ,这样在调用 func 时会自动启用推理模式,禁用梯度计算。# torch.is_inference_mode_enabled() 提供了一种方便的方式来检查当前是否处于推理模式,这在调试和性能优化时非常有用。# 一个条件判断语句。 TORCH_1_9 是一个变量,它用于判断PyTorch版本是否大于等于1.9.0。 torch.is_inference_mode_enabled() 是一个函数,用于检查当前是否已经启用了推理模式。如果这两个条件都为真,即PyTorch版本大于等于1.9.0且当前已经处于推理模式,那么执行下一行代码。if TORCH_1_9 and torch.is_inference_mode_enabled():# 如果上述条件判断为真,说明当前已经处于推理模式,无需再应用装饰器,因此直接返回原函数 fn ,相当于一个通过操作。return fn # already in inference_mode, act as a pass-through# 如果上述条件判断为假,即PyTorch版本低于1.9.0,或者虽然版本大于等于1.9.0但当前未启用推理模式,那么执行下一行代码。else:# 这行代码根据 TORCH_1_9 变量的值来选择装饰器。如果 TORCH_1_9 为真,则选择 torch.inference_mode 装饰器;否则选择 torch.no_grad 装饰器。然后立即调用所选装饰器,并将 fn 作为参数传递给装饰器,最终返回装饰后的函数。return (torch.inference_mode if TORCH_1_9 else torch.no_grad)()(fn)# 返回内部定义的 decorate 函数。这样,当 smart_inference_mode 被用作装饰器时,实际上返回的是 decorate 函数, decorate 函数会根据PyTorch的版本号来决定应用哪个装饰器。return decorate
# 这段代码通过定义一个装饰器工厂函数 smart_inference_mode ,根据PyTorch的版本号动态选择合适的装饰器来优化模型的推理模式。如果PyTorch版本大于等于1.9.0且当前已启用推理模式,则直接返回原函数;否则根据版本号选择 torch.inference_mode 或 torch.no_grad 装饰器来装饰函数,以提高推理效率。
4.def get_cpu_info():
# 这段代码定义了一个名为 get_cpu_info 的函数,用于获取当前计算机的 CPU 信息。它使用了 py-cpuinfo 库来获取 CPU 的详细信息。
# 定义一个名为 get_cpu_info 的函数,该函数不接受任何参数。
def get_cpu_info():# 返回包含系统 CPU 信息的字符串,例如‘Apple M2’。"""Return a string with system CPU information, i.e. 'Apple M2'."""# 导入 cpuinfo 模块,该模块用于获取 CPU 的详细信息。注释中提到需要通过 pip install py-cpuinfo 来安装这个模块。import cpuinfo # pip install py-cpuinfo# 定义一个元组 k ,包含三个字符串 "brand_raw" , "hardware_raw" , "arch_string_raw" 。这些字符串是 cpuinfo 模块返回的字典中可能包含的键,按照优先级排序(优先级从左到右递减)。注释说明不是所有的键在每次调用时都可用。# "brand_raw" :CPU 的品牌信息,通常包含完整的品牌名称和型号。 "hardware_raw" :硬件信息,可能包含一些硬件相关的描述。 "arch_string_raw" :架构信息,描述 CPU 的架构类型。k = "brand_raw", "hardware_raw", "arch_string_raw" # info keys sorted by preference (not all keys always available)# 调用 cpuinfo.get_cpu_info() 函数获取 CPU 信息,并将返回的字典赋值给变量 info 。info = cpuinfo.get_cpu_info() # info dict# 从 info 字典中尝试依次获取 k 元组中的键对应的值 :首先尝试获取 k[0] ,即 "brand_raw" 键对应的值; 如果 "brand_raw" 键不存在,则尝试获取 k[1] ,即 "hardware_raw" 键对应的值; 如果 "hardware_raw" 键也不存在,则尝试获取 k[2] ,即 "arch_string_raw" 键对应的值。# 如果以上键都不存在,则返回默认值 "unknown" 。将获取到的值赋值给变量 string 。string = info.get(k[0] if k[0] in info else k[1] if k[1] in info else k[2], "unknown")# 对变量 string 进行字符串替换操作。# 将字符串中的 "(R)" 替换为空字符串,即移除 "(R)" ; 将字符串中的 "CPU " 替换为空字符串,即移除 "CPU " ; 将字符串中的 "@ " 替换为空字符串,即移除 "@ " 。最后返回处理后的字符串。return string.replace("(R)", "").replace("CPU ", "").replace("@ ", "")
# 该函数 get_cpu_info 通过 cpuinfo 模块获取 CPU 的相关信息,并按照一定的优先级顺序尝试获取特定的键值。获取到的字符串值会经过一些简单的替换处理,以去除一些不需要的字符,最终返回一个简洁的 CPU 信息字符串。如果无法获取到有效的 CPU 信息,则返回 "unknown" 。
5.def select_device(device="", batch=0, newline=False, verbose=True):
# 这段代码定义了一个名为 select_device 的函数,用于选择和配置设备(如 CPU、GPU 或 MPS),以便在 PyTorch 框架中运行模型。
# 定义了一个函数 select_device ,接受四个参数。
# 1.device :指定要使用的设备,可以是字符串或 torch.device 对象。
# 2.batch :批处理大小,用于检查 GPU 设备数量与批处理大小是否匹配。
# 3.newline :布尔值,用于控制输出信息是否换行。
# 4.verbose :布尔值,用于控制是否打印设备选择信息。
def select_device(device="", batch=0, newline=False, verbose=True):# 根据提供的参数选择适当的 PyTorch 设备。# 该函数接受指定设备的字符串或 torch.device 对象,并返回表示所选设备的 torch.device 对象。该函数还会验证可用设备的数量,如果请求的设备不可用,则会引发异常。# 引发:# ValueError:如果指定的设备不可用,或者使用多个 GPU 时批处理大小不是设备数量的倍数。# 注意:# 设置“CUDA_VISIBLE_DEVICES”环境变量以指定要使用的 GPU。"""Selects the appropriate PyTorch device based on the provided arguments.The function takes a string specifying the device or a torch.device object and returns a torch.device objectrepresenting the selected device. The function also validates the number of available devices and raises anexception if the requested device(s) are not available.Args:device (str | torch.device, optional): Device string or torch.device object.Options are 'None', 'cpu', or 'cuda', or '0' or '0,1,2,3'. Defaults to an empty string, which auto-selectsthe first available GPU, or CPU if no GPU is available.batch (int, optional): Batch size being used in your model. Defaults to 0.newline (bool, optional): If True, adds a newline at the end of the log string. Defaults to False.verbose (bool, optional): If True, logs the device information. Defaults to True.Returns:(torch.device): Selected device.Raises:ValueError: If the specified device is not available or if the batch size is not a multiple of the number ofdevices when using multiple GPUs.Examples:>>> select_device('cuda:0')device(type='cuda', index=0)>>> select_device('cpu')device(type='cpu')Note:Sets the 'CUDA_VISIBLE_DEVICES' environment variable for specifying which GPUs to use."""# 如果 device 参数已经是一个 torch.device 对象。if isinstance(device, torch.device):# 直接返回该对象,无需进一步处理。return device# 这段代码是 select_device 函数的一部分,主要负责处理设备选择逻辑,确保用户请求的设备是可用的,并且配置正确。# 创建一个字符串 s ,用于记录当前环境的信息,包括 YOLO 的版本号、Python 的版本号和 PyTorch 的版本号。这个字符串用于后续的错误信息输出或日志记录。s = f"Ultralytics YOLOv{__version__} 🚀 Python-{PYTHON_VERSION} torch-{torch.__version__} "# 将 device 参数转换为字符串,并转换为小写,以便后续处理时能够统一处理不同的设备标识符格式。device = str(device).lower()# 遍历 remove 列表中的字符串。for remove in "cuda:", "none", "(", ")", "[", "]", "'", " ":# 将 device 中的这些字符串替换为空字符串。这一步的目的是简化设备标识符的格式,例如将 "cuda:0" 转换为 "0",将 "(0, 1)" 转换为 "0,1",以便后续更方便地处理设备标识符。device = device.replace(remove, "") # to string, 'cuda:0' -> '0' and '(0, 1)' -> '0,1'# 判断 device 是否为 "cpu",如果是,则将 cpu 变量设置为 True ,表示用户请求使用 CPU 设备。cpu = device == "cpu"# 判断 device 是否为 "mps" 或 "mps:0",如果是,则将 mps 变量设置为 True ,表示用户请求使用 Apple Metal Performance Shaders (MPS) 设备。mps = device in ("mps", "mps:0") # Apple Metal Performance Shaders (MPS)# 如果用户请求使用 CPU 或 MPS 设备,则设置环境变量 CUDA_VISIBLE_DEVICES 为 "-1"。这个操作会使得 torch.cuda.is_available() 返回 False ,从而确保 PyTorch 不会尝试使用 CUDA 设备。if cpu or mps:os.environ["CUDA_VISIBLE_DEVICES"] = "-1" # force torch.cuda.is_available() = False# 如果用户请求了非 CPU 设备(即请求了 GPU 设备)。elif device: # non-cpu device requested# 如果 device 为 "cuda",则将其设置为 "0",表示默认使用第一个 CUDA 设备。if device == "cuda":device = "0"# 获取当前环境变量 CUDA_VISIBLE_DEVICES 的值,并将其存储在 visible 变量中。visible = os.environ.get("CUDA_VISIBLE_DEVICES", None)# 设置环境变量 CUDA_VISIBLE_DEVICES 为 device ,以便后续操作可以识别指定的 CUDA 设备。这一步必须在调用 torch.cuda.is_available() 之前完成,以确保环境变量的设置生效。os.environ["CUDA_VISIBLE_DEVICES"] = device # set environment variable - must be before assert is_available()# 检查 CUDA 设备是否可用,并且设备数量是否满足用户请求的设备数量。如果 CUDA 设备不可用或设备数量不足,则执行后续的错误处理代码。if not (torch.cuda.is_available() and torch.cuda.device_count() >= len(device.split(","))):# 如果 CUDA 设备不可用或设备数量不足,则打印之前创建的字符串 s ,记录当前的版本信息。LOGGER.info(s)# 根据 CUDA 设备的数量,创建一个提示字符串 install 。如果 PyTorch 没有检测到任何 CUDA 设备,则提供 PyTorch 安装指南的链接,以便用户检查和安装 CUDA 设备驱动和相关依赖。install = ("See https://pytorch.org/get-started/locally/ for up-to-date torch install instructions if no ""CUDA devices are seen by torch.\n" # 如果 torch 没有看到 CUDA 设备,请参阅https://pytorch.org/get-started/locally/ 获取最新的 torch 安装说明。if torch.cuda.device_count() == 0else "")# 抛出 ValueError 异常,提示用户请求的 CUDA 设备无效,并建议用户使用 CPU 或提供有效的 CUDA 设备标识符。异常信息中包含了当前的 CUDA 设备状态和环境变量配置情况,以及之前创建的提示字符串 install ,以便用户进行问题排查和解决。raise ValueError(f"Invalid CUDA 'device={device}' requested." # 请求的 CUDA‘device={device}’无效。f" Use 'device=cpu' or pass valid CUDA device(s) if available," # 使用‘device=cpu’或传递有效的 CUDA 设备(如果可用),f" i.e. 'device=0' or 'device=0,1,2,3' for Multi-GPU.\n" # 即,对于多 GPU,‘device=0’或‘device=0,1,2,3’。f"\ntorch.cuda.is_available(): {torch.cuda.is_available()}"f"\ntorch.cuda.device_count(): {torch.cuda.device_count()}"f"\nos.environ['CUDA_VISIBLE_DEVICES']: {visible}\n"f"{install}")# 这段代码的主要作用是根据用户提供的设备参数和系统环境,进行设备选择和配置的初步处理。它会检查用户请求的设备类型,并设置相应的环境变量,以确保后续的 PyTorch 操作能够正确地使用指定的设备。同时,它还会对 CUDA 设备的可用性和数量进行检查,如果不符合要求,则抛出异常提示用户进行调整。# 这段代码是 select_device 函数中的一部分,主要负责根据系统的可用设备和配置选择合适的计算设备。# 如果当前不是CPU模式且不是MPS模式,并且CUDA设备可用,则优先选择GPU。if not cpu and not mps and torch.cuda.is_available(): # prefer GPU if available# 将设备字符串按逗号分割成列表,如果没有指定设备,则默认使用 "0" ,即第一个GPUdevices = device.split(",") if device else "0" # range(torch.cuda.device_count()) # i.e. 0,1,6,7# 获取设备数量。n = len(devices) # device count# 如果设备数量大于1且批量大小不为0且不能被设备数量整除,则抛出错误。if n > 1 and batch > 0 and batch % n != 0: # check batch_size is divisible by device_count# 错误信息提示用户调整批量大小,使其可以被设备数量整除,例如使用 batch // n * n 或 batch // n * n + n 来找到最接近的可整除的批量大小。raise ValueError(f"'batch={batch}' must be a multiple of GPU count {n}. Try 'batch={batch // n * n}' or " # batch={batch}' 必须是 GPU 数量 {n} 的倍数。尝试 'batch={batch // n * n}' 或f"'batch={batch // n * n + n}', the nearest batch sizes evenly divisible by {n}." # 'batch={batch // n * n + n}',最接近可以被 {n} 整除的批次大小。)# 计算用于对齐的空格数量。space = " " * (len(s) + 1)# 遍历每个设备编号。for i, d in enumerate(devices):# 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 设备的详细信息,例如显存大小、计算能力等,这些信息对于深度学习模型的训练和优化非常有用。# 获取当前设备的属性信息。p = torch.cuda.get_device_properties(i)# 将设备信息添加到字符串中,包括 设备名称 和 总内存大小 (以MB为单位)s += f"{'' if i == 0 else space}CUDA:{d} ({p.name}, {p.total_memory / (1 << 20):.0f}MiB)\n" # bytes to MB# 设置默认的设备参数为 "cuda:0" ,即第一个GPU。arg = "cuda:0"# 如果当前是MPS模式且PyTorch版本为2.0且MPS设备可用,则选择MPS。elif mps and TORCH_2_0 and torch.backends.mps.is_available():# Prefer MPS if available# 将MPS设备信息添加到字符串中。# def get_cpu_info(): -> 用于获取当前计算机的 CPU 信息。 -> return string.replace("(R)", "").replace("CPU ", "").replace("@ ", "")s += f"MPS ({get_cpu_info()})\n"# 设置设备参数为 "mps" 。arg = "mps"# 如果以上条件都不满足,则选择CPU。else: # revert to CPU# 将CPU设备信息添加到字符串中。s += f"CPU ({get_cpu_info()})\n"# 设置设备参数为 "cpu" 。arg = "cpu"# 这段代码通过检查系统的可用设备和配置,选择合适的计算设备,并构建相应的设备信息字符串。它优先选择GPU,其次是MPS,最后是CPU. 如果使用多GPU,还会检查批量大小是否可以均匀分配到每个GPU上,以避免潜在的性能问题。# 如果 verbose 参数为 True ,则根据 newline 参数的值,打印设备选择信息 s ,如果 newline 为 True ,则保留换行符,否则去除换行符。if verbose:LOGGER.info(s if newline else s.rstrip())# 返回一个 torch.device 对象,表示最终选择的设备。return torch.device(arg)
# 这个函数的主要作用是根据用户提供的设备参数和系统环境,选择合适的设备用于 PyTorch 模型的运行。它会优先选择 GPU 或 MPS 设备,如果不可用或配置不正确,则回退到 CPU 设备。同时,它还会检查批处理大小与设备数量的匹配情况,并在必要时抛出异常提示用户调整参数。
6.def time_sync():
# 这段代码定义了一个名为 time_sync 的函数,它用于在 PyTorch 中获取准确的时间。
# 这行定义了 time_sync 函数,它没有接受任何参数。
def time_sync():# PyTorch 精确的时间。"""PyTorch-accurate time."""# 检查 CUDA(Compute Unified Device Architecture,统一计算架构)是否可用,即是否有 NVIDIA 的 GPU 支持。if torch.cuda.is_available():# 如果 CUDA 可用,调用 torch.cuda.synchronize() 函数。这个函数用于确保所有的 CUDA 核心都完成了当前的任务,也就是说,它会等待所有先前排队的 CUDA 操作完成。# 这是必要的,因为在多线程环境中,不同线程可能会同时执行 CUDA 操作,而 time.time() 函数可能在这些操作完成之前就被调用,导致时间测量不准确。torch.cuda.synchronize()# 返回当前的时间,以秒为单位。 time.time() 函数返回自 Unix 纪元(1970年1月1日)以来的秒数。return time.time()
# time_sync 函数用于在 PyTorch 中获取一个同步的时间戳,特别是在使用 GPU 时,它确保了时间测量的准确性。在进行性能分析时,这个函数非常有用,因为它可以减少由于 GPU 异步操作导致的测量误差。
7.def fuse_conv_and_bn(conv, bn):
# 这段代码定义了一个名为 fuse_conv_and_bn 的函数,其目的是将一个卷积层( conv )和一个批归一化层( bn )融合成一个单独的卷积层( fusedconv )。
# 定义了 fuse_conv_and_bn 函数,它接受两个参数。
# 1.conv :一个卷积层。
# 2.bn :一个批归一化层 。
def fuse_conv_and_bn(conv, bn):# 融合 Conv2d() 和 BatchNorm2d() 层 https://tehnokv.com/posts/fusing-batchnorm-and-conv/。"""Fuse Conv2d() and BatchNorm2d() layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/."""# 创建一个新的 nn.Conv2d 卷积层 fusedconv ,其参数与原始卷积层 conv 相同。 fusedconv = (nn.Conv2d(conv.in_channels,conv.out_channels,kernel_size=conv.kernel_size,stride=conv.stride,padding=conv.padding,dilation=conv.dilation,groups=conv.groups,# 增加了 bias=True 参数,因为批归一化层的参数将被融合到卷积层的权重和偏置中。 bias=True,)# requires_grad_(False) 禁用了梯度计算,因为融合后的卷积层的权重不应该在训练中更新。 .requires_grad_(False)# to(conv.weight.device) 确保新卷积层的权重位于与原始卷积层相同的设备上。.to(conv.weight.device))# 这段代码是 fuse_conv_and_bn 函数中用于准备融合后的卷积层滤波器(权重)的部分。# Prepare filters# 首先克隆了原始卷积层 conv 的权重,以避免直接修改原始权重。# view(conv.out_channels, -1) 将权重重塑为一个二维张量,其中 : conv.out_channels 是卷积层输出通道的数量。 -1 表示自动计算剩余的维度,这样权重张量的每一行对应一个输出通道的所有卷积核权重。w_conv = conv.weight.clone().view(conv.out_channels, -1)# 创建了一个对角矩阵 w_bn ,其对角线上的元素是批归一化层 bn 的权重除以方差的平方根加上一个小的常数 eps (这是批归一化中用于防止除以零的技巧)。# torch.diag 函数用于从给定的一维张量创建对角矩阵。# bn.weight 是批归一化层的 可学习权重 ,它们用于缩放标准化后的值。# bn.running_var 是批归一化层在训练过程中计算的 运行方差 。# bn.eps 是一个小的常数,用于数值稳定性。w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var)))# 这行代码执行了以下步骤 :# torch.mm(w_bn, w_conv) :执行矩阵乘法,将批归一化层的权重(以对角矩阵形式)与卷积层的权重相乘。这个操作将批归一化的效果融合到卷积层的权重中。# .view(fusedconv.weight.shape) :将矩阵乘法的结果重塑为新卷积层 fusedconv 的权重形状。# fusedconv.weight.copy_(...) :将融合后的权重复制到新卷积层 fusedconv 的权重中。fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.shape))# 这部分代码的目的是将批归一化层的参数融合到卷积层的权重中。这是通过首先将卷积层的权重重塑为二维张量,然后创建一个包含批归一化权重的对角矩阵,最后执行矩阵乘法并将结果复制到新卷积层的权重中来实现的。这样,新卷积层的权重就包含了原始卷积层和批归一化层的效果。# 这段代码是 fuse_conv_and_bn 函数中用于准备融合后的卷积层空间偏置(bias)的部分。# Prepare spatial bias# 检查原始卷积层 conv 是否有偏置项。# 如果 conv 没有偏置项( conv.bias is None ),则创建一个形状为 conv.weight.shape[0] (即输出通道数)的全零张量作为偏置,并确保这个张量位于与 conv.weight 相同的设备上。# 如果 conv 有偏置项,则直接使用 conv.bias 。b_conv = torch.zeros(conv.weight.shape[0], device=conv.weight.device) if conv.bias is None else conv.bias# 计算 批归一化层的偏置项 ,考虑了均值和方差。# bn.bias 是批归一化层的 偏置项 。# bn.weight 是批归一化层的 可学习权重 ,用于缩放标准化后的值。# bn.running_mean 是批归一化层在训练过程中计算的 运行均值 。# bn.running_var 是批归一化层在训练过程中计算的 运行方差 。# bn.eps 是一个小的常数,用于数值稳定性。# 这个表达式计算了批归一化层的偏置项,减去了由于均值和方差导致的偏移。b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps))# 这行代码执行了以下步骤 :# b_conv.reshape(-1, 1) :将卷积层的偏置项重塑为列向量。# torch.mm(w_bn, b_conv.reshape(-1, 1)) :执行矩阵乘法,将 批归一化层的权重 (以对角矩阵形式)与 卷积层的偏置 向量相乘。这个操作将批归一化的效果融合到卷积层的偏置中。# .reshape(-1) :将矩阵乘法的结果重塑为一维张量。# + b_bn :将融合后的偏置与批归一化层的偏置项相加。# fusedconv.bias.copy_(...) :将融合后的偏置复制到新卷积层 fusedconv 的偏置中。fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn)# 这部分代码的目的是将批归一化层的参数融合到卷积层的偏置中。这是通过首先处理卷积层的偏置项(如果有的话),然后计算批归一化层的偏置项,接着执行矩阵乘法并将结果与批归一化层的偏置项相加,最后将融合后的偏置复制到新卷积层的偏置中来实现的。这样,新卷积层的偏置就包含了原始卷积层和批归一化层的效果。# 返回融合后的卷积层。return fusedconv
# fuse_conv_and_bn 函数的作用是将卷积层和批归一化层的参数融合,创建一个新的卷积层,其权重和偏置包含了原始卷积层和批归一化层的效果。这个融合过程通常用于模型部署,以减少模型的参数数量和计算量,同时保持模型性能。融合后的卷积层可以替代原始的卷积层和批归一化层,使得模型更加高效。
8.def fuse_deconv_and_bn(deconv, bn):
# 这段代码定义了一个名为 fuse_deconv_and_bn 的函数,其作用是将 ConvTranspose2d (反卷积)层和 BatchNorm2d (批量归一化)层融合为一个 ConvTranspose2d 层。这种融合可以减少模型的参数数量,提高模型的推理速度。
# 这行代码定义了一个名为 fuse_deconv_and_bn 的函数,它接受两个参数。
# 1.deconv 和 2.bn :分别代表 ConvTranspose2d 层和 BatchNorm2d 层。
def fuse_deconv_and_bn(deconv, bn):# 融合 ConvTranspose2d() 和 BatchNorm2d() 层。"""Fuse ConvTranspose2d() and BatchNorm2d() layers."""# 创建了一个新的 ConvTranspose2d 层 fuseddconv ,其参数与输入的 deconv 层相同,但将 bias 设置为 True ,表示该层将包含偏置项。然后,将该层的梯度计算设置为 False (即不参与梯度计算),并将其移动到与 deconv 层相同的设备上。fuseddconv = (nn.ConvTranspose2d(deconv.in_channels,deconv.out_channels,kernel_size=deconv.kernel_size,stride=deconv.stride,padding=deconv.padding,output_padding=deconv.output_padding,dilation=deconv.dilation,groups=deconv.groups,bias=True,).requires_grad_(False).to(deconv.weight.device))# 这三行代码是 fuse_deconv_and_bn 函数中用于融合 ConvTranspose2d 层和 BatchNorm2d 层权重的核心部分。# Prepare filters# 准备 ConvTranspose2d 层的权重,以便进行融合操作。# deconv.weight.clone() : deconv.weight 获取 ConvTranspose2d 层的权重张量。 clone() 方法创建该权重张量的一个副本,确保不会修改原始权重。# .view(deconv.out_channels, -1) : view 方法用于重塑张量的形状。 deconv.out_channels 是输出通道数。 -1 表示自动计算该维度的大小,使得重塑后的张量总元素数不变。# 例如,如果 deconv.weight 的形状是 (out_channels, in_channels, kernel_height, kernel_width) ,那么重塑后的 w_deconv 的形状将是 (out_channels, in_channels * kernel_height * kernel_width) 。w_deconv = deconv.weight.clone().view(deconv.out_channels, -1)# 准备 BatchNorm2d 层的权重,以便进行融合操作。# bn.weight.div(torch.sqrt(bn.eps + bn.running_var)) : bn.weight 获取 BatchNorm2d 层的权重张量。 bn.running_var 是运行时的方差张量。 bn.eps 是一个小的常数,用于数值稳定性。 torch.sqrt(bn.eps + bn.running_var) 计算方差加 eps 的平方根。 div 方法将 bn.weight 除以上述计算的平方根,得到归一化后的权重。# torch.diag(...) : torch.diag 方法创建一个对角矩阵,对角线上的元素是输入张量的元素。# 例如,如果输入张量是一个一维张量 [a, b, c] ,那么 torch.diag 将创建一个对角矩阵 [[a, 0, 0], [0, b, 0], [0, 0, c]] 。w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var)))# 将融合后的权重复制到新的 ConvTranspose2d 层 fuseddconv 中。# torch.mm(w_bn, w_deconv) : torch.mm 方法进行矩阵乘法。 w_bn 是一个对角矩阵,表示 BatchNorm2d 层的归一化权重。 w_deconv 是 ConvTranspose2d 层的权重,已经重塑为二维张量。 矩阵乘法的结果是一个新的二维张量,表示融合后的权重。# .view(fuseddconv.weight.shape) : view 方法将矩阵乘法的结果重塑为 fuseddconv.weight 的原始形状。# 例如,如果 fuseddconv.weight 的形状是 (out_channels, in_channels, kernel_height, kernel_width) ,那么重塑后的张量将具有相同的形状。# fuseddconv.weight.copy_(...) : copy_ 方法将计算得到的融合权重复制到 fuseddconv.weight 中,更新 fuseddconv 层的权重。fuseddconv.weight.copy_(torch.mm(w_bn, w_deconv).view(fuseddconv.weight.shape))# 这三行代码的核心作用是将 ConvTranspose2d 层的权重和 BatchNorm2d 层的权重进行融合,生成一个新的权重张量,并将其复制到新的 ConvTranspose2d 层 fuseddconv 中。重塑 ConvTranspose2d 层的权重为二维张量。计算 BatchNorm2d 层的归一化权重,并创建对角矩阵。通过矩阵乘法将两个权重张量相乘,得到融合后的权重。将融合后的权重重塑为原始形状,并复制到新的 ConvTranspose2d 层中。这种融合操作可以减少模型的参数数量,提高模型的推理速度。# 这三行代码是 fuse_deconv_and_bn 函数中用于融合 ConvTranspose2d 层和 BatchNorm2d 层偏置的核心部分。# Prepare spatial bias# 准备 ConvTranspose2d 层的偏置,以便进行融合操作。# deconv.weight.shape[1] : deconv.weight 获取 ConvTranspose2d 层的权重张量。 deconv.weight.shape[1] 是输入通道数。# torch.zeros(deconv.weight.shape[1], device=deconv.weight.device) : torch.zeros 方法创建一个全零张量,形状为 (deconv.weight.shape[1],) ,即输入通道数。 device=deconv.weight.device 确保新创建的张量与 deconv.weight 在同一个设备上(CPU或GPU)。# if deconv.bias is None else deconv.bias : 如果 deconv 层没有偏置项(即 deconv.bias 为 None ),则使用全零张量 b_conv 。 否则,使用 deconv 层的偏置 deconv.bias 。b_conv = torch.zeros(deconv.weight.shape[1], device=deconv.weight.device) if deconv.bias is None else deconv.bias# 计算 BatchNorm2d 层的偏置,以便进行融合操作。# bn.bias : bn.bias 获取 BatchNorm2d 层的偏置张量。# bn.weight.mul(bn.running_mean) : bn.weight 获取 BatchNorm2d 层的权重张量。 bn.running_mean 是运行时的均值张量。 mul 方法将 bn.weight 和 bn.running_mean 相乘。# div(torch.sqrt(bn.running_var + bn.eps)) : bn.running_var 是运行时的方差张量。 bn.eps 是一个小的常数,用于数值稳定性。 torch.sqrt(bn.running_var + bn.eps) 计算方差加 eps 的平方根。 div 方法将上述乘积结果除以平方根,得到归一化后的均值。# bn.bias - ... :最终, bn.bias 减去归一化后的均值,得到 BatchNorm2d 层的偏置 b_bn 。b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps))# 将融合后的偏置复制到新的 ConvTranspose2d 层 fuseddconv 中。# b_conv.reshape(-1, 1) : b_conv 是 ConvTranspose2d 层的偏置张量。 reshape(-1, 1) 将 b_conv 重塑为二维张量,形状为 (in_channels, 1) 。# torch.mm(w_bn, b_conv.reshape(-1, 1)) : w_bn 是 BatchNorm2d 层的对角矩阵。 torch.mm 方法进行矩阵乘法。 矩阵乘法的结果是一个新的二维张量,形状为 (out_channels, 1) 。# .reshape(-1) : reshape(-1) 将结果重塑为一维张量,形状为 (out_channels,) 。# + b_bn :将上述结果与 BatchNorm2d 层的偏置 b_bn 相加,得到融合后的偏置。# fuseddconv.bias.copy_(...) : copy_ 方法将计算得到的融合偏置复制到 fuseddconv.bias 中,更新 fuseddconv 层的偏置。fuseddconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn)# 这三行代码的核心作用是将 ConvTranspose2d 层的偏置和 BatchNorm2d 层的偏置进行融合,生成一个新的偏置张量,并将其复制到新的 ConvTranspose2d 层 fuseddconv 中。准备 ConvTranspose2d 层的偏置,如果该层没有偏置则使用全零张量。计算 BatchNorm2d 层的偏置,通过归一化均值和偏置相减得到。通过矩阵乘法将 BatchNorm2d 层的对角矩阵和 ConvTranspose2d 层的偏置相乘,然后加上 BatchNorm2d 层的偏置,得到融合后的偏置。将融合后的偏置复制到新的 ConvTranspose2d 层中。这种融合操作可以减少模型的参数数量,提高模型的推理速度。# 返回融合后的 ConvTranspose2d 层 fuseddconv 。return fuseddconv
# 这段代码通过定义一个函数 fuse_deconv_and_bn ,将 ConvTranspose2d 层和 BatchNorm2d 层融合为一个 ConvTranspose2d 层。创建一个新的 ConvTranspose2d 层 fuseddconv ,其参数与输入的 deconv 层相同,但包含偏置项。计算融合后的权重,通过矩阵乘法将 BatchNorm2d 层的权重和 ConvTranspose2d 层的权重相乘。计算融合后的偏置,通过矩阵乘法将 BatchNorm2d 层的权重和 ConvTranspose2d 层的偏置相乘,并加上 BatchNorm2d 层的偏置。返回融合后的 ConvTranspose2d 层 fuseddconv 。这种融合可以减少模型的参数数量,提高模型的推理速度。
9.def model_info(model, detailed=False, verbose=True, imgsz=640):
# 这段代码定义了一个名为 model_info 的函数,它用于打印和返回模型的详细信息,包括层数、参数数量、梯度数量以及计算复杂度(FLOPs)。
# 这行定义了 model_info 函数,它接受以下参数 :
# 1.model :要分析的模型。
# 2.detailed (默认为 False ) :一个布尔值,指示是否打印详细的参数信息。
# 3.verbose (默认为 True ) :一个布尔值,指示是否打印模型信息。
# 4.imgsz (默认为 640) :一个整数,表示模型输入图像的尺寸。
def model_info(model, detailed=False, verbose=True, imgsz=640):# 模型信息。# imgsz 可以是 int 或 list,即 imgsz=640 或 imgsz=[640, 320]。"""Model information.imgsz may be int or list, i.e. imgsz=640 or imgsz=[640, 320]."""# 如果 verbose 参数为 False ,则函数直接返回,不打印任何信息。if not verbose:return# 分别计算模型的 参数总数 、 梯度总数 和 层数 。n_p = get_num_params(model) # number of parametersn_g = get_num_gradients(model) # number of gradientsn_l = len(list(model.modules())) # number of layers# 检查 detailed 参数是否为 True ,如果是,则打印详细的参数信息。if detailed:# 打印详细的表头信息。LOGGER.info(f"{'layer':>5} {'name':>40} {'gradient':>9} {'parameters':>12} {'shape':>20} {'mu':>10} {'sigma':>10}")# 使用 enumerate 函数遍历模型的所有参数。 model.named_parameters() 返回一个迭代器,包含模型中所有参数的 名称 和 参数对象 。 enumerate 函数为每个参数添加了一个索引 i 。for i, (name, p) in enumerate(model.named_parameters()):# 移除参数名称中的前缀 "module_list." ,以便于阅读。这通常用于模型中包含嵌套模块时,简化参数名称。name = name.replace("module_list.", "")# 开始记录一条信息,使用 LOGGER 对象,这是一个常见的做法来记录重要的事件或潜在的问题。LOGGER.info(# 使用字符串格式化来构建一条包含参数详细信息的日志信息,并使用 LOGGER.info 打印出来。格式化字符串 %5g %40s %9s %12g %20s %10.3g %10.3g %10s 指定了每个字段的宽度和格式 :# %5g :索引 i ,占用5个字符宽度。# %40s :参数名称 name ,占用40个字符宽度。# %9s :参数是否需要梯度 p.requires_grad ,占用9个字符宽度。# %12g :参数元素总数 p.numel() ,占用12个字符宽度。# %20s :参数形状 list(p.shape) ,占用20个字符宽度。# %10.3g :参数均值 p.mean() ,占用10个字符宽度并保留3位小数。# %10.3g :参数标准差 p.std() ,占用10个字符宽度并保留3位小数。# %10s :参数数据类型 p.dtype ,占用10个字符宽度。"%5g %40s %9s %12g %20s %10.3g %10.3g %10s"% (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std(), p.dtype))# 计算模型的 FLOPs(浮点运算次数)。# def get_flops(model, imgsz=640): -> 用于计算 YOLO 模型的浮点运算次数(FLOPs)。 -> return thop.profile(deepcopy(model), inputs=[im], verbose=False)[0] / 1e9 * 2 # imgsz GFLOPsflops = get_flops(model, imgsz)# 检查模型是否已经融合,并根据结果添加相应的后缀。fused = " (fused)" if getattr(model, "is_fused", lambda: False)() else ""# 构造包含 FLOPs 信息的字符串。fs = f", {flops:.1f} GFLOPs" if flops else ""# 使用了 Python 的 getattr 函数来尝试从模型对象中获取 yaml_file 的值。 getattr 函数有三个参数 :对象、属性名和默认值。如果指定的属性在对象中不存在,它会返回默认值。# getattr(model, "yaml_file", "") :尝试获取模型对象 model 的 yaml_file 属性,如果不存在,则返回空字符串 "" 。# getattr(model, "yaml", {}) :如果 yaml_file 属性不存在,那么尝试获取模型对象 model 的 yaml 属性。如果 yaml 属性也不存在,它将返回一个空字典 {} 。# .get("yaml_file", "") :在获取到的 yaml 属性(可能是一个字典)上调用 get 方法,尝试获取键 "yaml_file" 对应的值。如果 "yaml_file" 键不存在, get 方法将返回空字符串 "" 。# or :这是一个逻辑或操作符。如果 getattr(model, "yaml_file", "") 返回的不是空字符串(即 yaml_file 属性存在且有值),那么 yaml_file 变量将被赋值为这个非空字符串。# 如果 getattr(model, "yaml_file", "") 返回空字符串,那么 or 操作符将评估 getattr(model, "yaml", {}).get("yaml_file", "") 表达式,并将其结果赋值给 yaml_file 变量。yaml_file = getattr(model, "yaml_file", "") or getattr(model, "yaml", {}).get("yaml_file", "")# 构造模型名称,如果 YAML 文件名存在,则从文件名中提取模型名称,否则使用默认名称 "Model"。model_name = Path(yaml_file).stem.replace("yolo", "YOLO") or "Model"# 印模型的总结信息,包括 模型名称 、 层数 、 参数数量 、 梯度数量 和 FLOPs 。LOGGER.info(f"{model_name} summary{fused}: {n_l} layers, {n_p} parameters, {n_g} gradients{fs}")# 函数返回模型的 层数 、 参数数量 、 梯度数量 和 FLOPs 。return n_l, n_p, n_g, flops
# model_info 函数提供了一个详细的模型分析工具,它可以打印模型的层数、参数数量、梯度数量和计算复杂度。如果 detailed 参数为 True ,它还会打印每个参数的详细信息。这个函数对于理解模型的结构和性能非常有用。
10.def get_num_params(model):
# 这段代码定义了一个名为 get_num_params 的函数,它用于计算模型中的参数总数。
# 这行定义了 get_num_params 函数,它接受一个参数。
# 1.model :即要计算参数总数的模型。
def get_num_params(model):# 返回 YOLO 模型中的参数总数。"""Return the total number of parameters in a YOLO model."""# 使用生成器表达式来计算模型中所有参数的元素总数。 model.parameters() 返回模型中所有参数的迭代器。对于每个参数 x , x.numel() 方法返回参数中元素的总数。 sum 函数将所有参数的元素数累加起来,返回模型中所有参数的总数。return sum(x.numel() for x in model.parameters())
# get_num_params 函数提供了一个简单的方式来获取模型中所有参数的总数,这对于了解模型的大小和复杂度非常有用。
11.def get_num_gradients(model):
# 这段代码定义了一个名为 get_num_gradients 的函数,它用于计算模型中需要梯度的参数的总数。
# 这行定义了 get_num_gradients 函数,它接受一个参数。
# 1.model :即要计算需要梯度的参数总数的模型。
def get_num_gradients(model):# 返回 YOLO 模型中具有梯度的参数总数。"""Return the total number of parameters with gradients in a YOLO model."""# 使用生成器表达式来计算模型中所有需要梯度的参数的元素总数。 model.parameters() 返回模型中所有参数的迭代器。对于每个参数 x ,如果 x.requires_grad 为 True ,则 x.numel() 方法返回参数中元素的总数。 sum 函数将所有需要梯度的参数的元素数累加起来,返回模型中所有需要梯度的参数的总数。return sum(x.numel() for x in model.parameters() if x.requires_grad)
# get_num_gradients 函数提供了一个简单的方式来获取模型中所有需要梯度的参数的总数,这对于了解模型的训练参数非常有用。在深度学习中,通常只有部分参数需要梯度,因为某些参数可能被固定或冻结,这个函数可以帮助我们了解哪些参数是可训练的。
12.def model_info_for_loggers(trainer):
# 这段代码定义了一个名为 model_info_for_loggers 的函数,其作用是根据传入的 trainer 对象,收集并返回模型的相关信息,用于日志记录。这些信息包括模型的参数数量、FLOPs(浮点运算次数)、推理时间等。具体实现取决于 trainer.args.profile 的值,即是否启用模型性能分析。
# 这行代码定义了一个名为 model_info_for_loggers 的函数,它接受一个参数。
# 1.trainer :这个参数是一个包含模型训练和验证信息的对象。
def model_info_for_loggers(trainer):# 返回包含有用模型信息的模型信息字典。"""Return model info dict with useful model information.Example:YOLOv8n info for loggers```pythonresults = {'model/parameters': 3151904,'model/GFLOPs': 8.746,'model/speed_ONNX(ms)': 41.244,'model/speed_TensorRT(ms)': 3.211,'model/speed_PyTorch(ms)': 18.755}```"""# 一个条件判断语句,检查 trainer.args.profile 是否为真。如果为真,表示需要进行模型性能分析,包括ONNX和TensorRT的推理时间。if trainer.args.profile: # profile ONNX and TensorRT times# 如果 trainer.args.profile 为真,从 ultralytics.utils.benchmarks 模块导入 ProfileModels 类。这个类用于分析模型的性能。from ultralytics.utils.benchmarks import ProfileModels# 创建一个 ProfileModels 对象,传入 trainer.last (最近训练的模型)和 trainer.device (模型所在的设备),调用 profile 方法进行性能分析。 profile 方法返回一个列表,取列表的第一个元素作为 results 。results = ProfileModels([trainer.last], device=trainer.device).profile()[0]# 从 results 字典中移除 "model/name" 键值对,因为这个信息在日志中可能不需要。results.pop("model/name")# 如果 trainer.args.profile 为假,表示不需要进行详细的性能分析,只返回PyTorch的推理时间。else: # only return PyTorch times from most recent validation# 创建一个字典 results ,包含模型的参数数量和FLOPs。results = {# get_num_params 函数用于获取模型的 参数数量 。"model/parameters": get_num_params(trainer.model),# get_flops 函数用于获取模型的 FLOPs ,并保留三位小数。"model/GFLOPs": round(get_flops(trainer.model), 3),}# 将 trainer.validator.speed["inference"] (最近验证的推理时间)添加到 results 字典中,并保留三位小数。这个值表示PyTorch模型的推理时间(单位:毫秒)。results["model/speed_PyTorch(ms)"] = round(trainer.validator.speed["inference"], 3)# 返回包含模型信息的字典 results 。return results
# 这段代码通过定义一个函数 model_info_for_loggers ,根据 trainer.args.profile 的值,收集并返回模型的相关信息,用于日志记录。如果 trainer.args.profile 为真,使用 ProfileModels 类进行模型性能分析,包括ONNX和TensorRT的推理时间,并移除不必要的信息。如果 trainer.args.profile 为假,只返回PyTorch的模型参数数量、FLOPs和最近验证的推理时间。将收集到的信息存储在字典 results 中,并返回该字典。这种设计可以根据不同的需求灵活地收集模型的性能信息,适用于不同的日志记录和性能分析场景。
13.def get_flops(model, imgsz=640):
# 这段代码定义了一个名为 get_flops 的函数,它用于计算 YOLO 模型的浮点运算次数(FLOPs)。
# 这行定义了 get_flops 函数,它接受两个参数。
# 1.model :要计算 FLOPs 的模型。
# 2.imgsz (默认为 640) :输入图像的尺寸。
def get_flops(model, imgsz=640):# 返回 YOLO 模型的 FLOP。"""Return a YOLO model's FLOPs."""# 检查 thop 库是否可用。 thop 是一个用于计算模型 FLOPs 和参数数量的库。如果 thop 不可用,函数返回 0.0 GFLOPs。if not thop:return 0.0 # if not installed return 0.0 GFLOPs# 尝试去除模型的 DataParallel 包装,以便 thop 可以正确分析模型。 try:# de_parallel 函数是一个自定义函数,用于从模型中移除 torch.nn.DataParallel 或 torch.nn.parallel.DistributedDataParallel 包装。model = de_parallel(model)# 获取模型的第一个参数,用于确定参数所在的设备(CPU 或 GPU)。p = next(model.parameters())# 确保 imgsz 是一个列表。如果 imgsz 是一个整数或浮点数,它将被扩展为一个包含两个元素的列表,这两个元素相同,代表图像的宽度和高度。if not isinstance(imgsz, list):imgsz = [imgsz, imgsz] # expand if int/float# 这部分代码是一个 try-except 块,它故意引发一个异常,以便跳过注释掉的代码块,直接执行 except 块中的代码。try:# Use stride size for input tensor# stride = max(int(model.stride.max()), 32) if hasattr(model, "stride") else 32 # max stride# im = torch.empty((1, p.shape[1], stride, stride), device=p.device) # input image in BCHW format# flops = thop.profile(deepcopy(model), inputs=[im], verbose=False)[0] / 1e9 * 2 # stride GFLOPs# return flops * imgsz[0] / stride * imgsz[1] / stride # imgsz GFLOPsraise Exception# 在 except 块中,代码创建一个空的输入张量 im ,其尺寸基于实际的图像尺寸 imgsz 。except Exception:# Use actual image size for input tensor (i.e. required for RTDETR models)im = torch.empty((1, p.shape[1], *imgsz), device=p.device) # input image in BCHW format# 然后使用 thop.profile 函数计算模型的 FLOPs。计算结果以 GFLOPs(十亿次浮点运算)为单位返回。return thop.profile(deepcopy(model), inputs=[im], verbose=False)[0] / 1e9 * 2 # imgsz GFLOPs# 如果在这个过程中发生任何异常。except Exception:# 函数将返回 0.0 。return 0.0
# get_flops 函数用于计算 YOLO 模型的 FLOPs,它考虑了模型的实际输入图像尺寸。如果 thop 库不可用,或者在计算过程中发生异常,函数将返回 0.0。这个函数对于评估模型的计算复杂度非常有用。
14.def get_flops_with_torch_profiler(model, imgsz=640):
# 这段代码定义了一个名为 get_flops_with_torch_profiler 的函数,其作用是使用PyTorch的内置性能分析器 torch.profiler 来计算模型的FLOPs(浮点运算次数)。这是一种替代 thop 库的方法,用于评估模型的计算复杂度。
# 这行代码定义了一个名为 get_flops_with_torch_profiler 的函数,它接受两个参数。
# 1.model :要评估的模型。
# 2.imgsz :输入图像的尺寸,默认为640。
def get_flops_with_torch_profiler(model, imgsz=640):# 计算模型 FLOP( thop 替代方案)。"""Compute model FLOPs (thop alternative)."""# 一个条件判断语句,检查是否使用了PyTorch 2.0或更高版本。 TORCH_2_0 用于表示PyTorch版本是否满足条件。if TORCH_2_0:# 如果条件为真,调用 de_parallel 函数处理模型,确保模型不是并行化的。这一步是为了确保模型在单个设备上运行,以便准确计算FLOPs。model = de_parallel(model)# 获取模型的第一个参数张量 p ,用于确定输入图像的通道数。p = next(model.parameters())# 计算模型的最大步长 stride 。如果模型有 stride 属性,则取其最大值并与32取最大值,然后乘以2。如果没有 stride 属性,则默认为64。这一步是为了确定输入图像的尺寸。stride = (max(int(model.stride.max()), 32) if hasattr(model, "stride") else 32) * 2 # max stride# 创建一个全零张量 im ,作为输入图像。张量的形状为 (1, p.shape[1], stride, stride) ,表示批量大小为1,通道数为 p.shape[1] ,高度和宽度为 stride 。张量的设备与模型参数 p 的设备相同。im = torch.zeros((1, p.shape[1], stride, stride), device=p.device) # input image in BCHW format# with torch.profiler.profile(activities, schedule, on_trace_ready, record_shapes, profile_memory, with_stack, with_flops, device, profile_memory_until, on_trace_ready_callback) as prof:# torch.profiler.profile() 是 PyTorch 提供的一个上下文管理器,用于性能分析,它可以测量和记录 PyTorch 操作的执行时间,以及 CPU 和 GPU 上的内存分配和释放事件。# 参数说明 :# activities :指定 profiler 需要监视的活动类型,可以是 ProfilerActivity.CPU 、 ProfilerActivity.CUDA 或它们的组合。# schedule :定义了 profiler 的调度策略,可以是 torch.profiler.schedule(wait, warmup, active, repeat) 。# wait :启动 profiler 前的等待时间。# warmup :预热时间,用于预热系统。# active :实际分析的时间。# repeat :重复分析的次数。# on_trace_ready :当跟踪准备好时的回调函数,可以用于将结果输出到文件或其他地方。# record_shapes :是否记录操作的输入和输出张量的形状。# profile_memory :是否分析内存分配和释放事件。# with_stack :是否记录操作的调用栈。# with_flops :是否计算浮点运算次数(FLOPs)。# device :指定要监视的设备,可以是 torch.device 或 torch.profiler.utils.get_device 返回的设备。# profile_memory_until :指定内存分析的持续时间。# on_trace_ready_callback :当跟踪准备好时的回调函数。# 返回值: :# prof :一个 Profile 对象,包含了分析的结果。# torch.profiler.profile() 提供了丰富的配置选项,允许用户根据具体需求定制性能分析的过程。通过合理配置这些参数,可以更有效地进行性能分析和优化。# 使用 torch.profiler.profile 上下文管理器进行性能分析,设置 with_flops=True 以启用FLOPs计算。在上下文管理器内部,调用模型 model 处理输入图像 im 。with torch.profiler.profile(with_flops=True) as prof:model(im)# 从性能分析器 prof 中获取 关键平均值 ,并计算总FLOPs。 x.flops 表示每个操作的FLOPs, sum 函数将所有操作的FLOPs相加,然后除以 1e9 将结果转换为GFLOPs(十亿浮点运算次数)。flops = sum(x.flops for x in prof.key_averages()) / 1e9# 如果 imgsz 是一个整数或浮点数,将其扩展为一个列表 [imgsz, imgsz] ,表示图像的宽度和高度。imgsz = imgsz if isinstance(imgsz, list) else [imgsz, imgsz] # expand if int/float# 根据输入图像的尺寸 imgsz 和步长 stride ,调整FLOPs的计算结果。这一步是为了将FLOPs从 stride x stride 的图像尺寸转换为 imgsz[0] x imgsz[1] 的图像尺寸。flops = flops * imgsz[0] / stride * imgsz[1] / stride # 640x640 GFLOPs# 返回计算得到的FLOPs。return flops# 如果条件判断为假(即PyTorch版本低于2.0),返回0。return 0
# 这段代码通过定义一个函数 get_flops_with_torch_profiler ,使用PyTorch的内置性能分析器 torch.profiler 来计算模型的FLOPs。检查PyTorch版本是否满足条件。处理模型,确保其不是并行化的。创建一个全零张量作为输入图像。使用 torch.profiler.profile 进行性能分析,计算FLOPs。根据输入图像的尺寸调整FLOPs的计算结果。返回计算得到的FLOPs。这种方法可以替代 thop 库,提供一种更灵活的FLOPs计算方式。
15.def initialize_weights(model):
# 这段代码定义了一个名为 initialize_weights 的函数,它用于初始化 PyTorch 模型中的权重和偏置。
# 定义 initialize_weights 函数,接受一个参数。
# 1.model :即需要初始化权重的 PyTorch 模型。
def initialize_weights(model):# 将模型权重初始化为随机值。"""Initialize model weights to random values."""# 遍历模型中的所有模块 m 。for m in model.modules():# 获取模块 m 的类型。t = type(m)# 如果模块 m 是 nn.Conv2d 类型,即二维卷积层,注释掉的代码表明通常这里会使用 He 初始化(也称为 Kaiming 初始化)来初始化权重。但由于被注释掉,所以实际上没有执行任何操作。if t is nn.Conv2d:pass # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')# 如果模块 m 是 nn.BatchNorm2d 类型,即二维批量归一化层,则设置其 eps 属性为 1e-3 (用于防止除以零的小值),并设置 momentum 属性为 0.03 (动量值)。elif t is nn.BatchNorm2d:m.eps = 1e-3m.momentum = 0.03# 如果模块 m 是激活函数中的某一种( nn.Hardswish 、 nn.LeakyReLU 、 nn.ReLU 、 nn.ReLU6 或 nn.SiLU ),则设置 inplace 属性为 True ,这允许激活函数在原地修改数据,减少内存消耗。elif t in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU]:m.inplace = True
# initialize_weights 函数负责初始化模型中的权重和偏置。在当前代码中,卷积层的权重初始化被跳过了,而批量归一化层的参数被设置了特定的值,激活函数被设置为原地操作。这种初始化对于模型的训练和性能是重要的,因为它可以帮助模型在训练初期更快地收敛,并减少梯度消失或爆炸的问题。
16.def scale_img(img, ratio=1.0, same_shape=False, gs=32):
# 这段代码定义了一个名为 scale_img 的函数,它用于对图像进行缩放操作,并可选地保持相同的形状或进行填充。
# 定义 scale_img 函数,接受以下参数 :
# 1.img :要缩放的图像。
# 2.ratio (默认为 1.0 ) :缩放比例。
# 3.same_shape (默认为 False ) :是否保持图像的原始形状。
# 4.gs (默认为 32 ) :网格尺寸,用于确定填充的边界。
def scale_img(img, ratio=1.0, same_shape=False, gs=32):# 根据给定的比例和网格大小 gs 缩放并填充形状为 img(bs,3,y,x) 的图像张量,可选择保留原始形状。"""Scales and pads an image tensor of shape img(bs,3,y,x) based on given ratio and grid size gs, optionallyretaining the original shape."""# 如果缩放比例为 1.0 ,即不进行缩放,直接返回原始图像。if ratio == 1.0:return img# 获取图像的高度和宽度。h, w = img.shape[2:]# 计算新的尺寸,即根据缩放比例计算出的新高度和宽度。s = (int(h * ratio), int(w * ratio)) # new size# 使用双线性插值方法对图像进行缩放, F.interpolate 是 PyTorch 中的插值函数。img = F.interpolate(img, size=s, mode="bilinear", align_corners=False) # resize# 如果 same_shape 为 False ,则计算新的填充尺寸,使得高度和宽度都是 gs 的整数倍。if not same_shape: # pad/crop imgh, w = (math.ceil(x * ratio / gs) * gs for x in (h, w))# 对图像进行填充,使得其尺寸符合新的填充尺寸。填充的值设置为 0.447 ,这是 ImageNet 数据集的均值之一。 F.pad 是 PyTorch 中的填充函数, [0, w - s[1], 0, h - s[0]] 指定了每个边界需要填充的像素数。return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447) # value = imagenet mean
# scale_img 函数用于根据给定的缩放比例对图像进行缩放,并可选地进行填充以保持特定的形状或尺寸。这个函数在图像预处理和数据增强中非常有用,尤其是在目标检测和图像分类任务中,它可以帮助模型更好地适应不同尺寸的输入图像。
17.def make_divisible(x, divisor):
# 这段代码定义了一个名为 make_divisible 的函数,其作用是将一个数 x 调整为最接近的、能被 divisor 整除的数。这个函数在深度学习中常用于调整卷积核大小、步长等参数,以确保它们符合特定的硬件要求或优化性能。
# 这行代码定义了一个名为 make_divisible 的函数,它接受两个参数。
# 1.x :要调整的数。
# 2.divisor :除数。
def make_divisible(x, divisor):# 返回最接近能被除数整除的 x。"""Returns nearest x divisible by divisor."""# 一个条件判断语句,检查 divisor 是否是一个 torch.Tensor 对象。if isinstance(divisor, torch.Tensor):# 如果 divisor 是一个 torch.Tensor ,则取其最大值并转换为整数。这一步确保 divisor 是一个整数,便于后续的计算。divisor = int(divisor.max()) # to int# 计算最接近 x 且能被 divisor 整除的数。# x / divisor :计算 x 除以 divisor 的商。# math.ceil(x / divisor) :使用 math.ceil 函数将商向上取整,确保结果是一个整数。# math.ceil(x / divisor) * divisor :将取整后的商乘以 divisor ,得到最接近 x 且能被 divisor 整除的数。return math.ceil(x / divisor) * divisor
# 这段代码通过定义一个函数 make_divisible ,将一个数 x 调整为最接近的、能被 divisor 整除的数。检查 divisor 是否是一个 torch.Tensor ,如果是,则取其最大值并转换为整数。计算 x 除以 divisor 的商,并向上取整。将取整后的商乘以 divisor ,得到最接近 x 且能被 divisor 整除的数。这个函数在深度学习中非常有用,特别是在设计卷积神经网络时,确保卷积核大小、步长等参数符合特定的硬件要求或优化性能。例如,某些硬件平台可能要求卷积核大小必须是特定值的倍数, make_divisible 函数可以方便地实现这一调整。
18.def copy_attr(a, b, include=(), exclude=()):
# 这段代码定义了一个名为 copy_attr 的函数,其作用是从对象 b 中复制属性到对象 a ,同时提供了包含( include )和排除( exclude )某些属性的选项。这个函数在对象属性管理中非常有用,特别是在需要部分复制对象属性的场景中。
# 这行代码定义了一个名为 copy_attr 的函数,它接受四个参数。
# 1.a :目标对象,属性将被复制到这个对象。
# 2.b :源对象,属性将从这个对象复制。
# 3.include :一个元组,包含需要复制的属性名称。默认为空元组,表示没有特定的属性需要包含。
# 4.exclude :一个元组,包含不需要复制的属性名称。默认为空元组,表示没有特定的属性需要排除。
def copy_attr(a, b, include=(), exclude=()):# 将属性从对象“b”复制到对象“a”,并可选择包含/排除某些属性。"""Copies attributes from object 'b' to object 'a', with options to include/exclude certain attributes."""# 遍历源对象 b 的所有属性。 b.__dict__ 是一个字典,包含对象 b 的所有属性及其值。 items() 方法返回一个包含键值对的迭代器。for k, v in b.__dict__.items():# 一个条件判断语句,用于决定是否跳过当前属性 k 。# len(include) and k not in include :如果 include 不为空且当前属性 k 不在 include 中,则跳过。# k.startswith("_") :如果属性名称以 _ 开头(通常表示私有属性),则跳过。# k in exclude :如果属性名称在 exclude 中,则跳过。if (len(include) and k not in include) or k.startswith("_") or k in exclude:continue# 如果当前属性 k 不满足上述任何跳过条件,则使用 setattr 函数将属性 k 及其值 v 设置到目标对象 a 中。else:setattr(a, k, v)
# 这段代码通过定义一个函数 copy_attr ,从对象 b 中复制属性到对象 a ,同时提供了包含和排除某些属性的选项。遍历源对象 b 的所有属性。根据 include 和 exclude 参数以及属性名称是否以 _ 开头,决定是否跳过当前属性。如果当前属性不满足跳过条件,则将该属性及其值复制到目标对象 a 中。这个函数在对象属性管理中非常有用,特别是在需要部分复制对象属性的场景中,例如在对象克隆、属性更新等操作中。通过灵活使用 include 和 exclude 参数,可以精确控制哪些属性需要被复制。
19.def get_latest_opset():
# 这段代码定义了一个名为 get_latest_opset 的函数,其作用是返回当前版本的PyTorch支持的次新ONNX opset版本号。ONNX opset是ONNX(Open Neural Network Exchange)定义的一组操作符(operators)的集合,每个opset版本号对应一组特定的操作符。这个函数通过检查PyTorch的 torch.onnx 模块中的符号操作集(symbolic opset)来确定次新版本号。
# 这行代码定义了一个名为 get_latest_opset 的函数,它不接受任何参数。
def get_latest_opset():# 返回此版本的 torch 最近支持的 ONNX opset(成熟度第二高)。"""Return second-most (for maturity) recently supported ONNX opset by this version of torch."""# vars(object)# vars() 函数在 Python 中用于获取对象的属性字典。这个字典包含了对象的大部分属性,但不包括方法和其他一些特殊的属性。对于用户自定义的对象, vars() 返回的字典包含了对象的 __dict__ 属性,这是一个包含对象所有属性的字典。# 参数说明 :# object :要获取属性字典的对象。# 返回值 :# 返回指定对象的属性字典。# 注意事项 :# vars() 对于内置类型(如 int 、 float 、 list 等)返回的是一个包含魔术方法和特殊属性的字典,这些属性通常是不可访问的。# 对于自定义对象, vars() 返回的是对象的 __dict__ 属性,如果对象没有定义 __dict__ ,则可能返回一个空字典或者抛出 TypeError 。# 在 Python 3 中, vars() 也可以用于获取内置函数的全局变量字典。# vars() 函数是一个内置函数,通常用于调试和访问对象的内部状态,但在处理复杂对象时应该谨慎使用,因为直接修改对象的属性可能会导致不可预测的行为。# 用于计算次新ONNX opset版本号。# vars(torch.onnx) : vars 函数返回一个对象的 __dict__ 属性,即一个包含对象所有属性和方法的字典。 torch.onnx 是PyTorch中与ONNX相关的模块, vars(torch.onnx) 返回该模块的所有属性和方法。# if "symbolic_opset" in k :这个条件判断语句用于筛选出包含 "symbolic_opset" 的属性名。这些属性名通常表示不同版本的ONNX opset的符号操作集。# int(k[14:]) : k[14:] 从属性名中提取版本号部分。假设属性名格式为 "symbolic_opset12" ,则 k[14:] 将提取出 "12" 。 int(k[14:]) 将提取出的字符串转换为整数。# max(int(k[14:]) for k in vars(torch.onnx) if "symbolic_opset" in k) :使用生成器表达式 int(k[14:]) for k in vars(torch.onnx) if "symbolic_opset" in k 生成所有opset版本号的列表。 max 函数返回这些版本号中的最大值,即最新支持的ONNX opset版本号。# - 1 :从最新版本号中减去1,得到次新版本号。return max(int(k[14:]) for k in vars(torch.onnx) if "symbolic_opset" in k) - 1 # opset
# 这段代码通过定义一个函数 get_latest_opset ,返回当前版本的PyTorch支持的次新ONNX opset版本号。获取 torch.onnx 模块的所有属性和方法。筛选出包含 "symbolic_opset" 的属性名。提取这些属性名中的版本号部分,并转换为整数。找到这些版本号中的最大值,即最新支持的ONNX opset版本号。从最新版本号中减去1,得到次新版本号。这个函数在需要确保使用成熟且稳定的ONNX opset版本时非常有用,特别是在将PyTorch模型导出为ONNX格式时。
20.def intersect_dicts(da, db, exclude=()):
# 这段代码定义了一个名为 intersect_dicts 的函数,它用于找出两个字典(通常用于存储模型的状态字典)中键和值形状都匹配的项。
# 这行定义了 intersect_dicts 函数,它接受三个参数。
# 1.da :第一个字典。
# 2.db :第二个字典。
# 3.exclude (默认为一个空元组) :一个元组,包含需要在比较时排除的键的部分。
def intersect_dicts(da, db, exclude=()):# 使用 da 值返回具有匹配形状的相交键的字典,不包括“exclude”键。"""Returns a dictionary of intersecting keys with matching shapes, excluding 'exclude' keys, using da values."""# 使用字典推导式来创建一个新的字典,其中只包含满足以下条件的键值对 :# k in db :键 k 必须同时存在于 da 和 db 中。# all(x not in k for x in exclude) :键 k 不能包含 exclude 元组中的任何字符串。# v.shape == db[k].shape :值 v 的形状必须与 db 中对应键的值的形状相匹配。return {k: v for k, v in da.items() if k in db and all(x not in k for x in exclude) and v.shape == db[k].shape}
# intersect_dicts 函数用于比较两个字典,找出它们共有的键,并且这些键的值的形状也相同。这个函数在深度学习中特别有用,例如在加载预训练模型时,可以确保只加载形状匹配的权重,避免因形状不匹配而导致的错误。通过排除特定的键,还可以进一步定制需要比较的字典项。
21.def is_parallel(model):
# 这段代码定义了一个名为 is_parallel 的函数,它用于检查一个 PyTorch 模型是否被 DataParallel 或 DistributedDataParallel 包装。
# 这行定义了 is_parallel 函数,它接受一个参数。
# 1.model :即要检查的 PyTorch 模型。
def is_parallel(model):# 如果模型是 DP 或 DDP 类型则返回 True。"""Returns True if model is of type DP or DDP."""# 使用 isinstance 函数来检查 model 是否是 nn.parallel.DataParallel 或 nn.parallel.DistributedDataParallel 类的实例。 isinstance 是 Python 的内建函数,用于检查一个对象是否是一个类或其子类的实例。# nn.parallel.DataParallel :这是 PyTorch 中用于数据并行处理的类,它可以在多个 GPU 上并行运行模型。# nn.parallel.DistributedDataParallel :这是 PyTorch 中用于分布式数据并行处理的类,它允许模型在多个节点上的多个 GPU 上并行运行。# 如果 model 是这两种类型中的任何一种, is_parallel 函数将返回 True ,表示模型被并行包装;否则,返回 False 。return isinstance(model, (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel))
# is_parallel 函数提供了一个简单的方式来检查 PyTorch 模型是否被 DataParallel 或 DistributedDataParallel 包装,这对于确定模型是否需要去除包装以访问原始模型结构非常有用。
22.def de_parallel(model):
# 这段代码定义了一个名为 de_parallel 的函数,它用于去除 PyTorch 模型的 DataParallel 或 DistributedDataParallel 包装,以获取原始模型。
# 这行定义了 de_parallel 函数,它接受一个参数。
# 1.model :即可能被 DataParallel 或 DistributedDataParallel 包装的 PyTorch 模型。
def de_parallel(model):# 对模型进行去并行化:如果模型类型为 DP 或 DDP,则返回单 GPU 模型。"""De-parallelize a model: returns single-GPU model if model is of type DP or DDP."""# 使用条件表达式来判断模型是否被 DataParallel 或 DistributedDataParallel 包装 :# is_parallel(model) :这它用于检查模型是否被 DataParallel 或 DistributedDataParallel 包装。如果模型被包装,该函数返回 True 。# model.module :如果模型被 DataParallel 或 DistributedDataParallel 包装, model.module 属性将指向原始模型。# else model :如果模型没有被包装,直接返回模型本身。return model.module if is_parallel(model) else model
# de_parallel 函数提供了一个简单的方式来获取 PyTorch 模型的原始版本,无论它是否被 DataParallel 或 DistributedDataParallel 包装。这对于需要直接访问模型内部结构的情况非常有用,例如在模型分析或保存模型时。
23.def one_cycle(y1=0.0, y2=1.0, steps=100):
# 这段代码定义了一个名为 one_cycle 的函数,其作用是生成一个“一次循环”(one-cycle)调度器。这种调度器常用于调整学习率或其他超参数,使其在一个训练周期内按照特定的曲线变化。具体来说,这个调度器生成的曲线是一个从 y1 到 y2 的平滑过渡,形状类似于一个半周期的余弦波。
# 这行代码定义了一个名为 one_cycle 的函数,它接受三个参数。
# 1.y1 :曲线的起始值,默认为0.0。
# 2.y2 :曲线的结束值,默认为1.0。
# 3.steps :训练周期的总步数,默认为100。
def one_cycle(y1=0.0, y2=1.0, steps=100):# 返回从 y1 到 y2 的正弦斜坡的 lambda 函数 https://arxiv.org/pdf/1812.01187.pdf。"""Returns a lambda function for sinusoidal ramp from y1 to y2 https://arxiv.org/pdf/1812.01187.pdf."""# 返回一个匿名函数(lambda函数),该函数接受一个参数 x ,表示当前的步数。这个lambda函数计算并返回当前步数对应的学习率或其他超参数值。# x * math.pi / steps :计算当前步数 x 在总步数 steps 中的比例,然后乘以 math.pi ,将步数映射到 [0, π] 区间。# math.cos(x * math.pi / steps) :计算上述值的余弦值,结果范围在 [-1, 1] 之间。# (1 - math.cos(x * math.pi / steps)) / 2 :将余弦值从 [-1, 1] 映射到 [0, 2] ,然后除以2,得到 [0, 1] 区间内的值。这个值表示从0到1的平滑过渡。# max((1 - math.cos(x * math.pi / steps)) / 2, 0) :使用 max 函数确保结果不小于0,虽然在这个映射范围内结果总是非负的,这一步主要是为了安全起见。# max((1 - math.cos(x * math.pi / steps)) / 2, 0) * (y2 - y1) :将上述结果乘以 (y2 - y1) ,将范围从 [0, 1] 扩展到 [0, y2 - y1] 。# max((1 - math.cos(x * math.pi / steps)) / 2, 0) * (y2 - y1) + y1 :最后,将结果加上 y1 ,将范围从 [0, y2 - y1] 平移到 [y1, y2] 。return lambda x: max((1 - math.cos(x * math.pi / steps)) / 2, 0) * (y2 - y1) + y1
# 这段代码通过定义一个函数 one_cycle ,生成一个“一次循环”调度器,用于在训练周期内平滑地调整学习率或其他超参数。定义起始值 y1 、结束值 y2 和总步数 steps 。返回一个lambda函数,该函数根据当前步数 x 计算并返回对应的学习率或其他超参数值。使用余弦函数生成从 y1 到 y2 的平滑过渡曲线。这种调度器在训练深度学习模型时非常有用,特别是在需要在训练周期内动态调整学习率以优化训练过程的场景中。通过这种平滑的过渡,可以避免学习率的突然变化,从而提高模型的收敛速度和稳定性。
24.def init_seeds(seed=0, deterministic=False):
# 这段代码定义了一个名为 init_seeds 的函数,其作用是初始化随机种子,以确保在使用Python的 random 模块、NumPy和PyTorch进行随机操作时能够复现结果。此外,该函数还提供了选项来控制是否启用确定性算法,这对于调试和复现实验结果非常有用。
# 这行代码定义了一个名为 init_seeds 的函数,它接受两个参数。
# 1.seed :随机种子的值,默认为0。
# 2.deterministic :一个布尔值,表示是否启用确定性算法,默认为 False 。
def init_seeds(seed=0, deterministic=False):# 初始化随机数生成器(RNG)种子 https://pytorch.org/docs/stable/notes/randomness.html。"""Initialize random number generator (RNG) seeds https://pytorch.org/docs/stable/notes/randomness.html."""# 使用 random.seed 函数设置Python标准库 random 模块的随机种子为 seed ,确保 random 模块生成的随机数序列是可复现的。random.seed(seed)# 使用 np.random.seed 函数设置NumPy的随机种子为 seed ,确保NumPy生成的随机数序列是可复现的。np.random.seed(seed)# 使用 torch.manual_seed 函数设置PyTorch的随机种子为 seed ,确保PyTorch生成的随机数序列是可复现的。torch.manual_seed(seed)# 使用 torch.cuda.manual_seed 函数设置PyTorch CUDA的随机种子为 seed ,确保在单个GPU上生成的随机数序列是可复现的。torch.cuda.manual_seed(seed)# 使用 torch.cuda.manual_seed_all 函数设置所有GPU的随机种子为 seed ,确保在多GPU环境下生成的随机数序列是可复现的。这个函数是异常安全的,即使在没有GPU的情况下也不会抛出异常。torch.cuda.manual_seed_all(seed) # for Multi-GPU, exception safe# 这行代码被注释掉了,如果取消注释, torch.backends.cudnn.benchmark 将被设置为 True 。这会启用CuDNN的自动调优功能,以选择最优的卷积算法,但可能会导致结果不可复现。# torch.backends.cudnn.benchmark = True # AutoBatch problem https://github.com/ultralytics/yolov5/issues/9287# 一个条件判断语句,检查 deterministic 是否为 True 。如果为 True ,则启用确定性算法。if deterministic:# 如果 deterministic 为 True ,进一步检查PyTorch版本是否为2.0或更高。 TORCH_2_0 用于表示PyTorch版本是否满足条件。if TORCH_2_0:# 如果PyTorch版本为2.0或更高,执行以下操作。# 启用PyTorch的确定性算法,并在无法启用时发出警告。torch.use_deterministic_algorithms(True, warn_only=True) # warn if deterministic is not possible# 启用CuDNN的确定性模式,确保卷积操作的结果是可复现的。torch.backends.cudnn.deterministic = True# 设置CuBLAS的工作空间配置,以确保线性代数操作的结果是可复现的。os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"# 设置Python的哈希种子,确保字符串哈希值是可复现的。os.environ["PYTHONHASHSEED"] = str(seed)# 如果PyTorch版本低于2.0,发出警告,建议用户升级到2.0或更高版本以启用确定性训练。else:LOGGER.warning("WARNING ⚠️ Upgrade to torch>=2.0.0 for deterministic training.") # 警告⚠️升级到 torch>=2.0.0 进行确定性训练。# 如果 deterministic 为 False ,则禁用确定性算法。else:# 禁用PyTorch的确定性算法。torch.use_deterministic_algorithms(False)# 禁用CuDNN的确定性模式。torch.backends.cudnn.deterministic = False
# 这段代码通过定义一个函数 init_seeds ,初始化随机种子,以确保在使用Python的 random 模块、NumPy和PyTorch进行随机操作时能够复现结果。此外,该函数还提供了选项来控制是否启用确定性算法,这对于调试和复现实验结果非常有用。设置Python标准库 random 模块的随机种子。设置NumPy的随机种子。设置PyTorch的随机种子。设置PyTorch CUDA的随机种子。设置所有GPU的随机种子。根据 deterministic 参数的值,启用或禁用确定性算法。如果启用确定性算法,根据PyTorch版本进行相应的配置。
25.class ModelEMA:
# 这段代码定义了一个名为 ModelEMA 的类,用于实现模型参数的指数移动平均(Exponential Moving Average, EMA)。EMA是一种常用的技巧,用于在训练过程中平滑模型参数,从而提高模型的泛化能力和稳定性。这个类提供了初始化EMA模型、更新EMA参数和更新模型属性的方法。
# 这行代码定义了一个名为 ModelEMA 的类。
class ModelEMA:# 更新了来自 https://github.com/rwightman/pytorch-image-models 的指数移动平均线 (EMA)# 保持模型 state_dict(参数和缓冲区)中所有内容的移动平均值# 有关 EMA 的详细信息,请参阅 https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage# 要禁用 EMA,请将 `enabled` 属性设置为 `False`。"""Updated Exponential Moving Average (EMA) from https://github.com/rwightman/pytorch-image-modelsKeeps a moving average of everything in the model state_dict (parameters and buffers)For EMA details see https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverageTo disable EMA set the `enabled` attribute to `False`."""# 这段代码是 ModelEMA 类的初始化方法 __init__ 的详细实现。这个方法用于创建和初始化一个用于模型参数指数移动平均(EMA)的对象。EMA是一种用于平滑模型参数的技术,有助于提高模型的泛化能力和稳定性。# 这行代码定义了 ModelEMA 类的初始化方法 __init__ ,它接受四个参数。# 1.model :要进行EMA的模型。# 2.decay :EMA的衰减率,默认为0.9999。# 3.tau :用于调整衰减率的参数,默认为2000。# 4.updates :EMA更新的次数,默认为0。def __init__(self, model, decay=0.9999, tau=2000, updates=0):"""Create EMA."""# 创建了一个EMA模型的副本,使用 deepcopy 确保模型参数完全独立。 de_parallel 函数用于处理模型的并行化,确保EMA模型不是并行化的。 eval 方法将模型设置为评估模式,确保在EMA更新过程中不会进行梯度计算。这一步确保EMA模型的参数不会参与到反向传播中,从而保持其作为平滑参数的作用。self.ema = deepcopy(de_parallel(model)).eval() # FP32 EMA# 将传入的 updates 值赋给 self.updates ,记录EMA更新的次数。self.updates = updates # number of EMA updates# 定义了一个lambda函数 self.decay ,用于计算当前更新次数下的衰减率。这个衰减率是一个指数衰减函数,帮助在训练初期更快地更新EMA模型。具体来说,随着更新次数的增加,衰减率逐渐接近 decay ,但在训练初期,衰减率会更快地增加,以便EMA模型能够更快地跟上训练模型的变化。self.decay = lambda x: decay * (1 - math.exp(-x / tau)) # decay exponential ramp (to help early epochs)# 遍历EMA模型的所有参数,并将它们的 requires_grad 属性设置为 False ,确保这些参数在训练过程中不会计算梯度。这一步是必要的,因为EMA模型的参数是通过平滑更新得到的,而不是通过反向传播更新的。for p in self.ema.parameters():p.requires_grad_(False)# 将 self.enabled 设置为 True ,表示EMA更新是启用的。这个属性可以用于在训练过程中动态启用或禁用EMA更新。self.enabled = True# 这段代码通过定义 ModelEMA 类的初始化方法 __init__ ,创建和初始化了一个用于模型参数指数移动平均(EMA)的对象。创建一个EMA模型的副本,确保模型参数完全独立,并设置为评估模式。记录EMA更新的次数。定义一个衰减率函数,用于在训练过程中调整EMA的更新速度。禁用EMA模型参数的梯度计算。启用EMA更新。这种EMA技术在训练深度学习模型时非常有用,特别是在需要提高模型泛化能力和稳定性的情况下。通过平滑模型参数,EMA可以减少模型在训练过程中的波动,从而提高最终模型的性能。# 这段代码是 ModelEMA 类的 update 方法的详细实现。这个方法用于在每次训练迭代后更新EMA模型的参数。EMA(指数移动平均)是一种用于平滑模型参数的技术,有助于提高模型的泛化能力和稳定性。# 这行代码定义了 ModelEMA 类的 update 方法,它接受一个参数。# 1.model :表示当前训练的模型。def update(self, model):"""Update EMA parameters."""# 一个条件判断语句,检查EMA更新是否启用。如果 self.enabled 为 True ,则继续执行更新操作。if self.enabled:# 如果EMA更新启用,更新次数 self.updates 加1。这一步记录了EMA模型已经进行了多少次更新。self.updates += 1# 计算当前更新次数下的衰减率 d 。 self.decay 是一个lambda函数,根据当前的更新次数 self.updates 计算衰减率。这个衰减率是一个指数衰减函数,帮助在训练初期更快地更新EMA模型。d = self.decay(self.updates)# 获取当前训练模型的状态字典 msd 。 de_parallel 函数用于处理模型的并行化,确保获取的状态字典是单个模型的状态字典,而不是并行模型的状态字典。msd = de_parallel(model).state_dict() # model state_dict# 遍历EMA模型的状态字典。 k 是参数的名称, v 是对应的参数值。for k, v in self.ema.state_dict().items():# 一个条件判断语句,检查参数 v 的数据类型是否为浮点数(FP16或FP32)。只有浮点数类型的参数才会进行EMA更新。if v.dtype.is_floating_point: # true for FP16 and FP32# 将当前EMA参数 v 乘以衰减率 d 。这一步是EMA更新的第一部分,即保留一部分旧的参数值。v *= d# 将 当前训练模型的参数 msd[k] 乘以 (1 - d) 并加到EMA参数 v 上。 detach 方法用于分离张量,避免计算图的扩展。这一步是EMA更新的第二部分,即引入一部分新的参数值。v += (1 - d) * msd[k].detach()# 一个断言语句,用于确保EMA参数 v 和训练模型参数 msd[k] 的数据类型都是 torch.float32 。虽然在实际使用中可能不需要这个断言,但它可以帮助调试和确保数据类型的一致性。# assert v.dtype == msd[k].dtype == torch.float32, f'{k}: EMA {v.dtype}, model {msd[k].dtype}'# 这段代码通过定义 ModelEMA 类的 update 方法,实现了在每次训练迭代后更新EMA模型的参数。检查EMA更新是否启用。更新EMA模型的更新次数。计算当前更新次数下的衰减率。获取当前训练模型的状态字典。遍历EMA模型的状态字典,更新每个浮点数类型的参数。使用衰减率和当前训练模型的参数,平滑更新EMA模型的参数。这种EMA技术在训练深度学习模型时非常有用,特别是在需要提高模型泛化能力和稳定性的情况下。通过平滑模型参数,EMA可以减少模型在训练过程中的波动,从而提高最终模型的性能。# 这段代码是 ModelEMA 类的 update_attr 方法的详细实现。这个方法用于更新EMA模型的属性,并保存去除优化器的模型。这在训练过程中非常有用,特别是在需要保持模型其他状态一致的情况下。# 这行代码定义了 ModelEMA 类的 update_attr 方法,它接受三个参数。# 1.model :当前训练的模型。# 2.include :一个元组,包含需要更新的属性名称,默认为空元组。# 3.exclude :一个元组,包含不需要更新的属性名称,默认为 ("process_group", "reducer") 。def update_attr(self, model, include=(), exclude=("process_group", "reducer")):# 更新属性并保存已删除优化器的剥离模型。"""Updates attributes and saves stripped model with optimizer removed."""# 一个条件判断语句,检查EMA更新是否启用。如果 self.enabled 为 True ,则继续执行更新操作。if self.enabled:# 如果EMA更新启用,调用 copy_attr 函数将当前训练模型的属性复制到EMA模型中。 copy_attr 函数会根据 include 和 exclude 参数进行筛选,确保只有需要的属性被更新。 include 参数用于指定需要更新的属性名称。 exclude 参数用于指定不需要更新的属性名称。# def copy_attr(a, b, include=(), exclude=()): -> 从对象 b 中复制属性到对象 a ,同时提供了包含( include )和排除( exclude )某些属性的选项。这个函数在对象属性管理中非常有用,特别是在需要部分复制对象属性的场景中。copy_attr(self.ema, model, include, exclude)# 这段代码通过定义 ModelEMA 类的 update_attr 方法,实现了更新EMA模型的属性,并保存去除优化器的模型。检查EMA更新是否启用。调用 copy_attr 函数,将当前训练模型的属性复制到EMA模型中,根据 include 和 exclude 参数进行筛选。这种属性更新机制在训练深度学习模型时非常有用,特别是在需要保持模型其他状态一致的情况下。通过更新EMA模型的属性,可以确保EMA模型在训练过程中与训练模型保持同步,从而提高最终模型的性能和稳定性。
# 这段代码通过定义一个 ModelEMA 类,实现了模型参数的指数移动平均(EMA)。初始化EMA模型,创建模型的副本并设置为评估模式。定义衰减率函数,用于在训练过程中调整EMA的更新速度。提供 update 方法,用于在每次训练迭代后更新EMA模型的参数。提供 update_attr 方法,用于更新EMA模型的属性,确保模型的其他状态也保持一致。这种EMA技术在训练深度学习模型时非常有用,特别是在需要提高模型泛化能力和稳定性的情况下。通过平滑模型参数,EMA可以减少模型在训练过程中的波动,从而提高最终模型的性能。
26.def strip_optimizer(f: Union[str, Path] = "best.pt", s: str = "") -> None:
# 这段代码定义了一个名为 strip_optimizer 的函数,其作用是从给定的PyTorch模型文件中移除优化器状态和其他不必要的信息,同时将模型转换为半精度(FP16)格式,并保存更新后的模型。这个函数特别适用于在模型训练完成后,准备模型用于推理或部署时,减少模型文件的大小并优化性能。
# 这行代码定义了 strip_optimizer 函数,它接受两个参数。
# 1.f :模型文件的路径,默认为 "best.pt" 。
# 2.s :保存更新后模型的路径,默认为空字符串,表示覆盖原文件。
def strip_optimizer(f: Union[str, Path] = "best.pt", s: str = "") -> None:# 从 'f' 中剥离优化器以完成训练,可选择保存为 's'。"""Strip optimizer from 'f' to finalize training, optionally save as 's'.Args:f (str): file path to model to strip the optimizer from. Default is 'best.pt'.s (str): file path to save the model with stripped optimizer to. If not provided, 'f' will be overwritten.Returns:NoneExample:```pythonfrom pathlib import Pathfrom ultralytics.utils.torch_utils import strip_optimizerfor f in Path('path/to/weights').rglob('*.pt'):strip_optimizer(f)```"""# 使用 torch.load 函数加载模型文件,并将所有张量映射到CPU上。 x 是一个包含模型信息的字典。x = torch.load(f, map_location=torch.device("cpu"))# 检查加载的模型字典中是否包含 "model" 键。如果不存在,输出一条信息并返回,表示这不是一个有效的Ultralytics模型。if "model" not in x:LOGGER.info(f"Skipping {f}, not a valid Ultralytics model.") # 跳过 {f},不是有效的 Ultralytics 模型。return# 如果模型对象有 args 属性,将其从 IterableSimpleNamespace 转换为字典。if hasattr(x["model"], "args"):x["model"].args = dict(x["model"].args) # convert from IterableSimpleNamespace to dict# 如果模型字典中包含 "train_args" ,将其与默认配置字典 DEFAULT_CFG_DICT 合并。否则, args 为 None 。args = {**DEFAULT_CFG_DICT, **x["train_args"]} if "train_args" in x else None # combine args# 如果模型字典中包含 "ema" (指数移动平均模型),用 "ema" 替换 "model" 。if x.get("ema"):x["model"] = x["ema"] # replace model with ema# 移除模型字典中的一些键,包括 "optimizer" 、 "best_fitness" 、 "ema" 和 "updates" ,将它们设置为 None 。for k in "optimizer", "best_fitness", "ema", "updates": # keysx[k] = None# 将模型的 "epoch" 设置为 -1 ,表示模型已经完成训练。x["epoch"] = -1# 将模型转换为半精度(FP16)格式,以减少模型大小并提高推理速度。x["model"].half() # to FP16# 遍历模型的所有参数,将 requires_grad 属性设置为 False ,表示这些参数在推理时不需要计算梯度。for p in x["model"].parameters():p.requires_grad = False# 更新 "train_args" ,只保留默认配置中的键。x["train_args"] = {k: v for k, v in args.items() if k in DEFAULT_CFG_KEYS} # strip non-default keys# x['model'].args = x['train_args']# 保存更新后的模型字典到指定路径。如果 s 为空字符串,则覆盖原文件。torch.save(x, s or f)# os.path.getsize(path)# os.path.getsize() 是 Python 标准库 os.path 模块中的一个函数,用于获取文件的大小,单位为字节。# 参数 :# path :文件的路径,可以是相对路径或绝对路径。# 返回值 :# 返回指定文件的大小,以字节为单位。如果文件不存在,会抛出 FileNotFoundError 异常。# 注意事项 :# os.path.getsize() 返回的是文件的实际大小,如果文件很大,返回的值可能非常大。# 这个函数只能用于获取文件的大小,不能用于目录。# 如果需要以更友好的方式显示文件大小(例如,自动转换为 KB、MB 或 GB),可能需要自己编写额外的代码来处理单位转换。# 计算保存后的模型文件大小(以MB为单位)。mb = os.path.getsize(s or f) / 1e6 # file size# 输出一条信息,表示优化器已从模型中移除,并显示保存路径和文件大小。LOGGER.info(f"Optimizer stripped from {f},{f' saved as {s},' if s else ''} {mb:.1f}MB") # 优化器从 {f} 中剥离,{f' 保存为 {s},' if s else ''} {mb:.1f}MB 。
# 这段代码通过定义 strip_optimizer 函数,实现了从PyTorch模型文件中移除优化器状态和其他不必要的信息,并将模型转换为半精度(FP16)格式,同时保存更新后的模型。这个函数特别适用于在模型训练完成后,准备模型用于推理或部署时,减少模型文件的大小并优化性能。加载模型文件。检查模型是否有效。转换模型的 args 属性。合并训练参数。用EMA模型替换原模型(如果存在)。移除不必要的键。将模型转换为FP16格式。禁用参数的梯度计算。保存更新后的模型。计算并输出模型文件的大小。
27.def profile(input, ops, n=10, device=None):
# 这段代码定义了一个名为 profile 的函数,用于分析和记录神经网络模型的性能指标,如参数数量、计算复杂度(GFLOPs)、内存使用情况、前向传播和反向传播的时间等。
# 定义一个名为 profile 的函数,接受以下参数 :
# 1.input :输入数据,可以是单个张量或张量列表。
# 2.ops :要分析的操作或模块,可以是单个操作或操作列表。
# 3.n :默认为10,表示重复执行操作的次数,用于计算平均时间。
# 4.device :指定设备,默认为 None ,表示自动选择设备。
def profile(input, ops, n=10, device=None):# Ultralytics 速度、内存和 FLOP 分析器。"""Ultralytics speed, memory and FLOPs profiler.Example:```pythonfrom ultralytics.utils.torch_utils import profileinput = torch.randn(16, 3, 640, 640)m1 = lambda x: x * torch.sigmoid(x)m2 = nn.SiLU()profile(input, [m1, m2], n=100) # profile over 100 iterations```"""# 这段代码是 profile 函数的开头部分,主要用于初始化结果列表、选择设备,并打印性能分析的表头信息。# 初始化一个空列表 results ,用于存储后续性能分析的结果。每个操作的分析结果将被添加到这个列表中,以便在函数结束时返回。results = []# 检查 device 参数是否是 torch.device 类型。如果不是,则调用 select_device 函数来选择一个合适的设备(例如CPU或GPU)。# select_device 函数通常会根据传入的 device 参数(如字符串 "cpu"、"0"、"cuda:0" 等)来决定使用哪个设备。如果 device 为 None 或未指定,则会自动选择一个可用的设备,通常是GPU(如果可用)。# 这样做的目的是确保在后续的操作中,输入数据和模型都能在正确的设备上运行,从而避免设备不匹配导致的错误。if not isinstance(device, torch.device):# 如果不是,则调用 select_device 函数来选择一个合适的设备(例如CPU或GPU)。# select_device 函数通常会根据传入的 device 参数(如字符串 "cpu"、"0"、"cuda:0" 等)来决定使用哪个设备。如果 device 为 None 或未指定,则会自动选择一个可用的设备,通常是GPU(如果可用)。# 这样做的目的是确保在后续的操作中,输入数据和模型都能在正确的设备上运行,从而避免设备不匹配导致的错误。# def select_device(device="", batch=0, newline=False, verbose=True): -> 用于选择和配置设备(如 CPU、GPU 或 MPS),以便在 PyTorch 框架中运行模型。返回一个 torch.device 对象,表示最终选择的设备。 -> return torch.device(arg)device = select_device(device)# 使用 LOGGER.info 打印性能分析的表头信息,定义了各个性能指标的列标题和对齐方式。# LOGGER 是一个日志记录器对象,通常用于记录程序运行时的信息、警告和错误等。LOGGER.info(# f-string 用于格式化字符串,其中 :>12s 、 :>14s 和 :>24s 表示右对齐,并分别指定列宽为12、14和24个字符。# Params :参数数量,表示模型或操作中的可训练参数总数。# GFLOPs :每秒十亿次浮点运算(Giga Floating-point Operations per second),用于衡量计算复杂度。# GPU_mem (GB) :GPU内存使用情况(以GB为单位),表示在执行操作时GPU的内存占用。# forward (ms) :前向传播时间(以毫秒为单位),表示执行一次前向传播所需的时间.。# backward (ms) :反向传播时间(以毫秒为单位),表示执行一次反向传播所需的时间。f"{'Params':>12s}{'GFLOPs':>12s}{'GPU_mem (GB)':>14s}{'forward (ms)':>14s}{'backward (ms)':>14s}"# input :输入形状,表示输入数据的形状或结构。# output :输出形状,表示操作执行后输出数据的形状或结构.f"{'input':>24s}{'output':>24s}")# 这段代码为后续的性能分析做好了准备,确保设备选择正确,并输出了清晰的表头信息,方便后续记录和查看每个操作的性能指标。# 这段代码是性能分析函数中的核心循环部分,主要负责对输入数据和操作进行处理,并计算操作的计算复杂度(GFLOPs)。# 一个条件表达式,用于遍历输入数据 input 。# 如果 input 是一个列表,则直接遍历列表中的每个元素 x 。 如果 input 不是列表,则将其视为单个输入,创建一个只包含 input 的列表 [input] ,然后遍历该列表。这种处理方式使得函数可以灵活地接受单个输入或多个输入。for x in input if isinstance(input, list) else [input]:# 将输入数据 x 移动到指定的设备 device 上,以便在该设备上进行后续的计算操作。x = x.to(device)# 设置输入数据 x 的 requires_grad 属性为 True ,表示需要对 x 进行梯度计算,以便在反向传播时计算梯度。x.requires_grad = True# 一个条件表达式,用于遍历操作 ops 。# 如果 ops 是一个列表,则直接遍历列表中的每个操作 m 。 如果 ops 不是列表,则将其视为单个操作,创建一个只包含 ops 的列表 [ops] ,然后遍历该列表。这种处理方式使得函数可以灵活地接受单个操作或多个操作。for m in ops if isinstance(ops, list) else [ops]:# 将操作 m 移动到指定的设备 device 上,以便在该设备上进行后续的计算操作。# hasattr(m, "to") 检查操作 m 是否具有 to 方法,通常 nn.Module 类型的操作具有 to 方法,可以将其移动到指定设备上。如果操作 m 不具有 to 方法,则不进行设备移动操作,直接使用原操作。m = m.to(device) if hasattr(m, "to") else m # device# 将操作 m 转换为半精度( half )。# hasattr(m, "half") 检查操作 m 是否具有 half 方法,通常 nn.Module 类型的操作具有 half 方法,可以将其转换为半精度。# isinstance(x, torch.Tensor) 检查输入数据 x 是否为 torch.Tensor 类型。# x.dtype is torch.float16 检查输入数据 x 的数据类型是否为 torch.float16 ,即半精度浮点数。# 如果以上条件都满足,则将操作 m 转换为半精度,否则不进行转换。# 使用半精度可以减少计算和存储需求,但可能会牺牲一定的精度。m = m.half() if hasattr(m, "half") and isinstance(x, torch.Tensor) and x.dtype is torch.float16 else m# 初始化时间变量 tf 和 tb ,分别用于记录 前向传播 和 反向传播 的时间, t 是一个列表,用于存储 时间戳 。tf, tb, t = 0, 0, [0, 0, 0] # dt forward, backwardtry:# 使用 thop.profile 函数计算操作 m 的计算复杂度(GFLOPs)。# thop.profile(m, inputs=[x], verbose=False) 调用 thop 库的 profile 函数,传入操作 m 和输入数据 [x] , verbose=False 表示不输出详细信息。# [0] 获取计算复杂度的值(以浮点运算次数为单位)。# / 1e9 将计算复杂度转换为每秒浮点运算次数(GFLOPs)。# * 2 表示将计算复杂度乘以2,通常用于表示双向计算(前向和反向传播)的复杂度。# if thop else 0 如果 thop 模块可用,则计算GFLOPs,否则设置为0.flops = thop.profile(m, inputs=[x], verbose=False)[0] / 1e9 * 2 if thop else 0 # GFLOPs# 使用 try-except 语句捕获可能出现的异常,如果计算GFLOPs时发生异常,则将 flops 设置为0,避免程序中断。except Exception:flops = 0# 这段代码通过处理输入数据和操作,并计算操作的计算复杂度,为后续的性能分析提供了基础数据。# 这段代码是性能分析函数中用于计算操作的前向和反向传播时间、内存使用情况、参数数量等指标,并记录和返回结果的部分。# 使用 try-except 语句捕获可能出现的异常,确保即使发生异常也能继续执行后续代码。try:# 循环 n 次,用于多次运行操作以计算平均时间。for _ in range(n):# 记录前向传播开始的时间戳。# def time_sync(): -> 用于在 PyTorch 中获取准确的时间。返回当前的时间,以秒为单位。 time.time() 函数返回自 Unix 纪元(1970年1月1日)以来的秒数。 -> return time.time()t[0] = time_sync()# 执行操作 m 的前向传播,将输入数据 x 传递给操作 m ,得到输出 y 。y = m(x)# 记录前向传播结束的时间戳。t[1] = time_sync()# 再次使用 try-except 语句捕获可能出现的异常,确保即使反向传播失败也能继续执行后续代码。try:# 执行反向传播。如果输出 y 是列表,则先对列表中的每个元素 yi 求和,然后对结果求和并调用 backward() 方法进行反向传播。如果输出 y 不是列表,则直接对 y 求和并调用 backward() 方法进行反向传播。(sum(yi.sum() for yi in y) if isinstance(y, list) else y).sum().backward()# 记录反向传播结束的时间戳。t[2] = time_sync()# 如果反向传播失败(例如操作没有反向传播方法),则将 t[2] 设置为 NaN (非数字),表示反向传播时间无效。except Exception: # no backward method# print(e) # for debugt[2] = float("nan")# 计算前向传播时间(以毫秒为单位),累加到 tf 中。tf += (t[1] - t[0]) * 1000 / n # ms per op forward# 计算反向传播时间(以毫秒为单位),累加到 tb 中。tb += (t[2] - t[1]) * 1000 / n # ms per op backward# 计算内存使用情况和参数数量。# 计算GPU内存使用情况(以GB为单位)。如果CUDA可用,则使用 torch.cuda.memory_reserved() 获取当前GPU上保留的内存大小,然后除以 1e9 转换为GB。如果CUDA不可用,则将内存使用量设置为0。mem = torch.cuda.memory_reserved() / 1e9 if torch.cuda.is_available() else 0 # (GB)# 获取 输入 和 输出 的 形状 。如果输入或输出是 torch.Tensor 类型,则使用 tuple(x.shape) 获取其形状。如果输入或输出不是 torch.Tensor 类型,则将其形状设置为 "list" 。s_in, s_out = (tuple(x.shape) if isinstance(x, torch.Tensor) else "list" for x in (x, y)) # shapes# 如果操作 m 是 nn.Module 类型,则遍历其参数并使用 x.numel() 获取每个参数的元素数量,然后求和得到总参数数量。如果操作 m 不是 nn.Module 类型,则将参数数量设置为0。p = sum(x.numel() for x in m.parameters()) if isinstance(m, nn.Module) else 0 # parameters# 记录和输出结果。# 使用日志记录器 LOGGER 的 info 方法输出操作的性能分析结果。使用 f-string 格式化字符串,将参数数量、GFLOPs、GPU内存使用、前向和反向传播时间、输入和输出形状等信息格式化输出。LOGGER.info(f"{p:12}{flops:12.4g}{mem:>14.3f}{tf:14.4g}{tb:14.4g}{str(s_in):>24s}{str(s_out):>24s}")# 将操作的性能分析结果添加到 results 列表中。results.append([p, flops, mem, tf, tb, s_in, s_out])# 如果在计算过程中发生异常,则记录异常信息并将结果设置为 None 。except Exception as e:LOGGER.info(e)results.append(None)# 清空GPU缓存,以避免内存泄露。torch.cuda.empty_cache()# 返回包含所有操作性能分析结果的列表 results ,包括 参数数量 、 GFLOPs 、 GPU内存使用 、 前向和反向传播时间 、 输入和输出形状 等信息格式化输出。return results# 这段代码通过多次运行操作来计算平均的前向和反向传播时间,同时记录内存使用情况和参数数量等指标,并将结果记录和返回,为后续的性能分析和优化提供了详细的数据支持.
# profile 函数通过多次迭代执行神经网络操作的前向传播和反向传播,计算并记录了操作的参数数量、计算复杂度(GFLOPs)、GPU内存使用情况、前向传播和反向传播的时间等性能指标。这些信息有助于分析和优化神经网络模型的性能,例如识别计算瓶颈、内存占用较大的操作等,从而指导模型的改进和优化工作。
28.class EarlyStopping:
# 这段代码定义了一个名为 EarlyStopping 的类,用于在训练过程中实现早停(Early Stopping)。早停是一种常用的技巧,用于在模型的性能在一定数量的训练周期内没有改善时停止训练,以防止过拟合。
# 这行代码定义了一个名为 EarlyStopping 的类。
class EarlyStopping:# 早期停止类,当经过指定数量的周期而没有得到改进时,停止训练。"""Early stopping class that stops training when a specified number of epochs have passed without improvement."""# 这段代码是 EarlyStopping 类的初始化方法 __init__ 的详细实现。这个方法用于初始化早停(Early Stopping)对象,设置一些初始状态和参数,以便在训练过程中监控模型的性能并决定是否提前停止训练。# 这行代码定义了 EarlyStopping 类的初始化方法 __init__ ,它接受一个参数。# 1.patience :在没有性能提升的情况下,等待的训练周期数,默认为50。def __init__(self, patience=50):# 初始化早期停止对象。"""Initialize early stopping object.Args:patience (int, optional): Number of epochs to wait after fitness stops improving before stopping."""# 初始化 最佳性能指标 (例如mAP)为0.0。这个值用于记录训练过程中模型的最佳性能。self.best_fitness = 0.0 # i.e. mAP# 初始化 最佳性能所在的训练周期 为0。这个值用于记录模型达到最佳性能的训练周期。self.best_epoch = 0# 设置 耐心值 patience ,表示在没有性能提升的情况下,等待的训练周期数。如果 patience 为0,则设置为无穷大,表示不使用早停。这一步确保了即使 patience 为0,也不会立即停止训练,而是继续训练直到手动停止。self.patience = patience or float("inf") # epochs to wait after fitness stops improving to stop# 初始化一个标志 possible_stop 为 False ,表示下一次训练周期是否可能停止训练。这个标志在训练过程中用于提前警告可能的停止,以便在下一个周期进行最终检查。self.possible_stop = False # possible stop may occur next epoch# 这段代码通过定义 EarlyStopping 类的初始化方法 __init__ ,初始化了一个早停对象,设置了一些初始状态和参数。初始化最佳性能指标为0.0。初始化最佳性能所在的训练周期为0。设置耐心值,表示在没有性能提升的情况下,等待的训练周期数。如果 patience 为0,则设置为无穷大。初始化一个标志,表示下一次训练周期是否可能停止训练。这种早停机制在训练深度学习模型时非常有用,特别是在需要防止过拟合和优化训练时间的情况下。通过合理设置耐心值,可以在模型性能不再提升时及时停止训练,节省计算资源并提高模型的泛化能力。# 这段代码是 EarlyStopping 类的 __call__ 方法的详细实现。这个方法用于在每个训练周期结束时调用,以检查是否应该提前停止训练。具体来说,它根据当前周期的性能指标(例如mAP)和最佳性能指标来决定是否停止训练。# 这行代码定义了 EarlyStopping 类的 __call__ 方法,它接受两个参数。# 1.epoch :当前训练周期。# 2.fitness :当前训练周期的性能指标(例如mAP)。def __call__(self, epoch, fitness):# 检查是否停止训练。"""Check whether to stop training.Args:epoch (int): Current epoch of trainingfitness (float): Fitness value of current epochReturns:(bool): True if training should stop, False otherwise"""# 如果性能指标 fitness 为 None ,表示当前周期没有进行验证,返回 False ,继续训练。if fitness is None: # check if fitness=None (happens when val=False)return False# 如果当前周期的性能指标 fitness 大于或等于最佳性能指标 self.best_fitness ,更新最佳性能指标和最佳周期。这一步确保了在训练初期,即使性能指标为0,也能正确更新最佳性能。if fitness >= self.best_fitness: # >= 0 to allow for early zero-fitness stage of trainingself.best_epoch = epochself.best_fitness = fitness# 计算自最佳性能以来的训练周期数。 delta 表示从最佳性能周期到当前周期的间隔。delta = epoch - self.best_epoch # epochs without improvement# 如果自最佳性能以来的训练周期数大于或等于 patience - 1 ,设置 possible_stop 为 True ,表示下一次训练周期可能停止训练。这一步用于提前警告可能的停止,以便在下一个周期进行最终检查。self.possible_stop = delta >= (self.patience - 1) # possible stop may occur next epoch# 如果自最佳性能以来的训练周期数大于或等于 patience ,设置 stop 为 True ,表示停止训练。这一步是早停的核心逻辑,确保在性能不再提升时停止训练。stop = delta >= self.patience # stop training if patience exceeded# 如果 stop 为 True ,输出一条信息,表示训练将提前停止,并提供最佳性能所在的周期和如何更新早停策略的建议。这一步用于记录和通知用户训练的停止原因和最佳模型的保存位置。if stop:LOGGER.info(f"Stopping training early as no improvement observed in last {self.patience} epochs. " # 由于在最后 {self.patience} 个时期未观察到任何改善,因此提前停止训练。f"Best results observed at epoch {self.best_epoch}, best model saved as best.pt.\n" # 在时期 {self.best_epoch} 观察到的最佳结果,最佳模型保存为 best.pt。f"To update EarlyStopping(patience={self.patience}) pass a new patience value, " # 要更新 EarlyStopping(patience={self.patience}),请传递一个新的耐心值,f"i.e. `patience=300` or use `patience=0` to disable EarlyStopping." # 即 `patience=300` 或使用 `patience=0` 来禁用 EarlyStopping。)# 返回 stop ,表示是否停止训练。如果 stop 为 True ,训练将停止;否则,训练将继续。return stop# 这段代码通过定义 EarlyStopping 类的 __call__ 方法,实现了在每个训练周期结束时检查是否应该提前停止训练。检查当前周期的性能指标是否为 None ,如果是,继续训练。如果当前周期的性能指标提升,更新最佳性能指标和最佳周期。计算自最佳性能以来的训练周期数。根据耐心值决定是否可能在下一个周期停止训练。根据耐心值决定是否停止训练。如果决定停止训练,输出相关信息并返回 True ,否则返回 False 。这种早停机制在训练深度学习模型时非常有用,特别是在需要防止过拟合和优化训练时间的情况下。通过合理设置耐心值,可以在模型性能不再提升时及时停止训练,节省计算资源并提高模型的泛化能力。
# 这段代码通过定义 EarlyStopping 类,实现了在训练过程中根据性能指标的提升情况决定是否提前停止训练。初始化最佳性能指标、最佳周期和耐心值。在每个训练周期调用 __call__ 方法,检查当前周期的性能指标。如果当前周期的性能指标提升,更新最佳性能指标和最佳周期。计算自最佳性能以来的训练周期数。根据耐心值决定是否停止训练。如果决定停止训练,输出相关信息并返回 True ,否则返回 False 。这种早停机制在训练深度学习模型时非常有用,特别是在需要防止过拟合和优化训练时间的情况下。通过合理设置耐心值,可以在模型性能不再提升时及时停止训练,节省计算资源并提高模型的泛化能力。