torch_utils.py
yolov6\utils\torch_utils.py
目录
torch_utils.py
1.所需的库和模块
2.def torch_distributed_zero_first(local_rank: int):
3.def time_sync():
4.def initialize_weights(model):
5.def fuse_conv_and_bn(conv, bn):
6.def fuse_model(model):
7.def get_model_info(model, img_size=640):
1.所需的库和模块
#!/usr/bin/env python3
# -*- coding:utf-8 -*-import time
from contextlib import contextmanager
from copy import deepcopy
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.nn.functional as F
from yolov6.utils.events import LOGGERtry:import thop # for FLOPs computation
except ImportError:thop = None2.def torch_distributed_zero_first(local_rank: int):
@contextmanager
def torch_distributed_zero_first(local_rank: int): # torch_distributed_zero_first torch分布式零优先# 装饰器使分布式训练中的所有进程等待每个 local_master 执行某项操作。"""Decorator to make all processes in distributed training wait for each local_master to do something."""if local_rank not in [-1, 0]:# dist.barrier()# 这是一个分布式同步操作,用于确保所有参与分布式训练的进程在继续执行之前都已经完成了当前的操作。# 当所有进程执行到这个 barrier 点时,它们会等待直到所有进程都到达这个 barrier 点。# 这样可以确保所有进程在进入下一个训练步骤之前都已经完成了当前步骤的工作。# 这个同步操作对于确保数据一致性和训练过程的正确性非常重要。dist.barrier(device_ids=[local_rank])yieldif local_rank == 0:dist.barrier(device_ids=[0])3.def time_sync():
def time_sync():# 如果 cuda 可用，则等待 CUDA 设备上所有流中的所有内核完成。'''Waits for all kernels in all streams on a CUDA device to complete if cuda is available.'''if torch.cuda.is_available():# torch.cuda.synchronize()# 是PyTorch库中的一个函数，它可以用于同步CPU和GPU之间的计算。这个函数会阻塞CPU线程，直到所有在当前设备上排队的CUDA核心完成执行为止。# 这样，当 torch.cuda.synchronize() 函数返回时，我们可以确保GPU的计算已经完成，并且数据已经准备好被CPU使用。# 通常情况下， torch.cuda.synchronize() 函数会在需要获取GPU计算结果的时候被使用，# 比如在训练深度神经网络时，需要在每个epoch结束后计算验证集的误差，此时就需要使用这个函数来同步CPU和GPU之间的计算，以确保得到正确的结果。torch.cuda.synchronize()return time.time()4.def initialize_weights(model):
# 它用于初始化传入的模型（ model ）中的权重和偏置。这个函数遍历模型中的所有模块，并根据模块的类型应用特定的初始化策略。
# model ：一个神经网络模型。
def initialize_weights(model):# 遍历模型中的所有模块。 model.modules() 返回一个迭代器，包含模型中的所有子模块。for m in model.modules():# 获取当前模块 m 的类型。t = type(m)if t is nn.Conv2d:# 不进行任何操作（ pass ）。这意味着卷积层的权重和偏置将保持它们默认的初始化状态，通常是由PyTorch自动初始化的。passelif t is nn.BatchNorm2d:m.eps = 1e-3m.momentum = 0.03elif t in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU]:# 将 inplace 参数设置为 True ，这意味着激活函数的计算将在输入数据上直接进行，而不是创建新的输出张量。这可以减少内存使用，但会修改输入数据。m.inplace = True
5.def fuse_conv_and_bn(conv, bn):
def fuse_conv_and_bn(conv, bn):# 融合卷积和批量规范层 https://tehnokv.com/posts/fusing-batchnorm-and-conv/。'''Fuse convolution and batchnorm layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/.'''fusedconv = (nn.Conv2d(conv.in_channels,conv.out_channels,kernel_size=conv.kernel_size,stride=conv.stride,padding=conv.padding,groups=conv.groups,bias=True,).requires_grad_(False).to(conv.weight.device))# prepare filters 准备过滤器# conv.weight 获取该卷积层的权重，这是一个包含卷积核参数的张量（tensor）。# .clone() 方法创建了 conv.weight 的一个副本。使用 .clone() 是为了避免在原始权重张量上直接进行操作，这样可以保留原始的权重数据不变。# conv.out_channels 是卷积层输出通道的数量，这里用作新张量的第一个维度大小。# .view(conv.out_channels, -1) 将权重张量重新塑形为一个二维张量，其中行数等于输出通道数，列数则是权重张量中剩余元素的总数除以输出通道数。# 这行代码的目的是将卷积层的权重张量重新塑形为一个二维张量，其中每一行代表一个卷积核的权重。这种操作通常用于分析卷积层的权重，或者在将卷积层和批量归一化层融合时计算新的卷积核权重。# 例如，如果一个卷积层有 64 个输出通道，每个卷积核的大小为 3x3，则原始权重张量的形状可能是 [64, 3, 3, 3] 。执行上述操作后， w_conv 的形状将变为 [64, 9] ，其中每一行代表一个卷积核的权重。w_conv = conv.weight.clone().view(conv.out_channels, -1)# torch.diag()# 函数可以用于创建对角线张量或提取给定矩阵的对角线元素。# 创建对角矩阵 ：如果输入是一个向量（1D张量）， torch.diag 将返回一个2D方阵，其中输入向量的元素作为对角线元素。# 提取对角线元素 ：如果输入是一个矩阵（2D张量）， torch.diag 将返回一个1D张量，包含输入矩阵的对角线元素。# 指定对角线 ： torch.diag 函数还允许你通过 diagonal 参数指定要提取或使用的对角线。 diagonal=0 表示主对角线， diagonal>0 表示主对角线上方的对角线， diagonal<0 表示主对角线下方的对角线。w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var)))# .copy_()# b = torch.empty_like(a).copy_(a)# copy_() 函数是需要一个目标tensor，也就是说需要先构建 b ，然后将 a 拷贝给 b ，而clone操作则不需要。# torch.mm(mat1, mat2) → Tensor# 函数返回一个新的张量，它是 mat1 和 mat2 的矩阵乘积。注意， mat1 的列数必须等于 mat2 的行数，以满足矩阵乘法的规则。# mat1 (Tensor) ：第一个要进行矩阵乘法的张量。通常是一个二维张量（矩阵）。# mat2 (Tensor) ：第二个要进行矩阵乘法的张量。也通常是一个二维张量（矩阵）。fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.shape))# prepare spatial bias 准备空间偏置# 如果卷积层没有偏置（ conv.bias 是 None ），则创建一个与卷积层输出通道数相同、设备相同的零张量作为偏置。# 如果卷积层有偏置，则直接使用 conv.bias 。b_conv = (torch.zeros(conv.weight.size(0), device=conv.weight.device)if conv.bias is Noneelse conv.bias)# torch.div()# div_(value) ：div() 的 in-place 运算形式。# .torch.div(input, value, out=None)# 将 input 逐元素除以标量值 value ，并返回结果到输出张量 out 。 即 out=tensor/value 。# 如果输入是 FloatTensor or DoubleTensor 类型，则参数 value 必须为实数，否则须为整数。# 参数：# input (Tensor) ：输入张量。# value (Number) ：除数。# out (Tensor, optional) ：输出张量。# bn.bias ： 是批量归一化层的偏置。# bn.weight ： 是批量归一化层的缩放因子。# bn.running_mean ： 是批量归一化层的运行均值。# bn.running_var ： 是批量归一化层的运行方差。# bn.eps ： 是批量归一化层的epsilon值，用于数值稳定性。# 这个表达式计算了在考虑缩放因子和均值/方差调整后的偏置。b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps))# w_bn ： 是批量归一化层的权重矩阵。# b_conv ：是卷积层的偏置，已经根据是否有偏置进行了初始化。# torch.mm(w_bn, b_conv.reshape(-1, 1)) ： 执行矩阵乘法，将权重矩阵 w_bn 与偏置向量 b_conv 相乘。# .reshape(-1) ： 将结果张量重新塑形为一维张量。# + b_bn ： 将批量归一化层调整后的偏置 b_bn 加到结果上。# fusedconv.bias.copy_(...) ： 将计算得到的融合偏置复制到融合后的卷积层的偏置张量中。fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn)return fusedconv6.def fuse_model(model):
# 这段代码定义了一个名为 fuse_model 的函数，其目的是对传入的模型进行优化，通过融合卷积层（Convolutional layers）和批量归一化层（Batch Normalization layers）来减少模型中的层数，提高推理速度。这种融合通常在模型部署或量化时进行，以优化模型性能。
# 定义了一个函数 fuse_model ，它接受一个参数 model ，这个参数是一个模型实例。
def fuse_model(model):# 融合模型的卷积层和批量规范层。'''Fuse convolution and batchnorm layers of the model.''' # 函数的文档字符串，说明了函数的作用。from yolov6.layers.common import ConvModule# model.modules()# 在PyTorch框架中， .modules() 方法是 nn.Module 类的一个成员函数，用于获取模型中所有的子模块。这个方法会递归地遍历模型，返回一个迭代器，包含模型本身及其所有子模块的引用。# 参数 ： 无参数。# 返回值 ： 返回一个迭代器，包含模型中的所有模块（包括模型自身）。# 遍历模型中的所有模块。 model.modules() 返回一个迭代器，包含模型中的所有子模块。for m in model.modules():# 检查当前模块 m 是否是 ConvModule 类型，并且是否具有批量归一化层 bn 属性。if type(m) is ConvModule and hasattr(m, "bn"):# 如果当前模块满足条件，则调用 fuse_conv_and_bn 函数来融合卷积层和批量归一化层，并将结果赋值给 m.conv 。m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv 更新 conv# 删除模块的 bn 属性，因为批量归一化层已经被融合到卷积层中。delattr(m, "bn") # remove batchnorm 删除 batchnorm# 将模块的 forward 方法替换为融合后的 forward_fuse 方法，这样在前向传播时会使用融合后的层。m.forward = m.forward_fuse # update forward 更新 forward# 返回修改后的模型。return model7.def get_model_info(model, img_size=640):
def get_model_info(model, img_size=640):# 获取模型参数和 GFlops。# 代码库位于https://github.com/Megvii-BaseDetection/YOLOX/blob/main/yolox/utils/model_utils.py"""Get model Params and GFlops.Code base on https://github.com/Megvii-BaseDetection/YOLOX/blob/main/yolox/utils/model_utils.py"""from thop import profilestride = 64 #32img = torch.zeros((1, 3, stride, stride), device=next(model.parameters()).device)# thop.profile()# 函数返回了两个值： flops 表示模型的计算量（浮点运算次数）， params 表示模型的参数量# copy.deepcopy()# 是深拷贝，会拷贝对象及其子对象，哪怕以后对其有改动，也不会影响其第一次的拷贝。# 函数可以递归地复制对象及其所有嵌套对象，创建一个全新的独立副本，而不共享任何数据。flops, params = profile(deepcopy(model), inputs=(img,), verbose=False)params /= 1e6flops /= 1e9img_size = img_size if isinstance(img_size, list) else [img_size, img_size]flops *= img_size[0] * img_size[1] / stride / stride * 2 # Gflops# 参数：{:.2f}M，Gflops：{:.2f}。info = "Params: {:.2f}M, Gflops: {:.2f}".format(params, flops)return info
