DETR纯代码分享（七）backbone.py(models)

一、导入模块/函数

from collections import OrderedDict

import torch
import torch.nn.functional as F
import torchvision
from torch import nn
from torchvision.models._utils import IntermediateLayerGetter
from typing import Dict, List

from util.misc import NestedTensor, is_main_process

from .position_encoding import build_position_encoding

这段代码导入了一些需要的Python模块和函数，以及一些自定义的模块和函数。以下是对每个导入的部分的简要解释：

from collections import OrderedDict：从Python的collections模块中导入OrderedDict类，这是一个有序字典，用于维护项目的插入顺序。
import torch：导入PyTorch深度学习框架的主要模块。
import torch.nn.functional as F：导入PyTorch的nn.functional模块，并将其命名为F，用于访问各种神经网络函数，如激活函数、损失函数等。
import torchvision：导入PyTorch的视觉工具库torchvision，它包含了用于计算机视觉任务的模型、数据集和变换等工具。
from torch import nn：从PyTorch的nn模块中导入nn类，这是构建神经网络模型的基础模块。
from torchvision.models._utils import IntermediateLayerGetter：从torchvision中的models._utils模块导入IntermediateLayerGetter类，它用于提取模型中间层的特征。
from typing import Dict, List：从Python的typing模块导入Dict和List，这些是用于类型提示的类，用于指定变量的类型。
from util.misc import NestedTensor, is_main_process：从自定义的util.misc模块中导入NestedTensor类和is_main_process函数。NestedTensor用于处理嵌套张量数据，而is_main_process函数用于检查当前进程是否为主进程。（具体实现可见misc.py文件）
from .position_encoding import build_position_encoding：从当前目录中的position_encoding模块导入build_position_encoding函数。这个函数用于构建位置编码（Position Encoding）。（具体实现可见position_encoding.py文件）

这些导入语句为后续的代码提供了所需的模块和函数，以构建和操作深度学习模型。

二、`FrozenBatchNorm2d模块`

class FrozenBatchNorm2d(torch.nn.Module):
    """
    BatchNorm2d where the batch statistics and the affine parameters are fixed.

    Copy-paste from torchvision.misc.ops with added eps before rqsrt,
    without which any other models than torchvision.models.resnet[18,34,50,101]
    produce nans.
    """

    def __init__(self, n):
        super(FrozenBatchNorm2d, self).__init__()
        self.register_buffer("weight", torch.ones(n))
        self.register_buffer("bias", torch.zeros(n))
        self.register_buffer("running_mean", torch.zeros(n))
        self.register_buffer("running_var", torch.ones(n))

    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
                              missing_keys, unexpected_keys, error_msgs):
        num_batches_tracked_key = prefix + 'num_batches_tracked'
        if num_batches_tracked_key in state_dict:
            del state_dict[num_batches_tracked_key]

        super(FrozenBatchNorm2d, self)._load_from_state_dict(
            state_dict, prefix, local_metadata, strict,
            missing_keys, unexpected_keys, error_msgs)

    def forward(self, x):
        # move reshapes to the beginning
        # to make it fuser-friendly
        w = self.weight.reshape(1, -1, 1, 1)
        b = self.bias.reshape(1, -1, 1, 1)
        rv = self.running_var.reshape(1, -1, 1, 1)
        rm = self.running_mean.reshape(1, -1, 1, 1)
        eps = 1e-5
        scale = w * (rv + eps).rsqrt()
        bias = b - rm * scale
        return x * scale + bias

这段代码定义了一个名为FrozenBatchNorm2d的自定义PyTorch模块，它是一个冻结的批量归一化（Batch Normalization）层，用于神经网络模型中。这个自定义模块实现了冻结的批量归一化操作，其中权重和偏差是固定的，而均值和方差是从运行中累积的，用于规范化输入数据。这对于某些模型（如ResNet系列）的性能很重要，以避免产生NaN值。

1、init()

class FrozenBatchNorm2d(torch.nn.Module):
    """
    BatchNorm2d where the batch statistics and the affine parameters are fixed.

    Copy-paste from torchvision.misc.ops with added eps before rqsrt,
    without which any other models than torchvision.models.resnet[18,34,50,101]
    produce nans.
    """

    def __init__(self, n):
        super(FrozenBatchNorm2d, self).__init__()
        self.register_buffer("weight", torch.ones(n))
        self.register_buffer("bias", torch.zeros(n))
        self.register_buffer("running_mean", torch.zeros(n))
        self.register_buffer("running_var", torch.ones(n))

这部分代码定义了FrozenBatchNorm2d类的构造函数__init__，它用于初始化一个FrozenBatchNorm2d的实例。这个类是用来实现冻结的批量归一化（Frozen Batch Normalization）操作的。下面是每行代码的详细解释：

def __init__(self, n):：这是构造函数的定义，它接受一个参数n，表示归一化的特征通道数。
super(FrozenBatchNorm2d, self).__init__()：通过调用父类torch.nn.Module的构造函数来初始化这个类的实例。
self.register_buffer("weight", torch.ones(n))：使用register_buffer方法注册一个名为weight的缓冲区（buffer），用于存储归一化的权重。这个权重的初始值被设置为一个全为1的张量，这意味着在冻结的批量归一化中，权重始终保持不变。
self.register_buffer("bias", torch.zeros(n))：同样，使用register_buffer方法注册一个名为bias的缓冲区，用于存储归一化的偏差。初始值为一个全为0的张量，这也是一个固定的值。
self.register_buffer("running_mean", torch.zeros(n))：注册一个名为running_mean的缓冲区，用于存储运行时的均值。初始值为一个全为0的张量，这是一个累积的值，在每个批次中更新。
self.register_buffer("running_var", torch.ones(n))：注册一个名为running_var的缓冲区，用于存储运行时的方差。初始值为一个全为1的张量，这也是一个累积的值，在每个批次中更新。

总之，这个构造函数初始化了一个冻结的批量归一化层，其中权重和偏差是固定的，而均值和方差是在运行时进行累积的。这种设置适用于某些神经网络模型，特别是那些在训练中使用批量归一化，并希望在推断过程中固定归一化参数的情况。

2、`_load_from_state_dict()`

    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
                              missing_keys, unexpected_keys, error_msgs):
        num_batches_tracked_key = prefix + 'num_batches_tracked'
        if num_batches_tracked_key in state_dict:
            del state_dict[num_batches_tracked_key]

        super(FrozenBatchNorm2d, self)._load_from_state_dict(
            state_dict, prefix, local_metadata, strict,
            missing_keys, unexpected_keys, error_msgs)

这部分代码定义了FrozenBatchNorm2d类的内部方法_load_from_state_dict，该方法用于加载模型的状态字典（state_dict）。以下是代码的详细解释：

def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):：这是方法的定义，它接受多个参数，包括模型的状态字典state_dict、前缀prefix等。
num_batches_tracked_key = prefix + 'num_batches_tracked'：创建一个变量num_batches_tracked_key，它是由参数prefix和字符串'num_batches_tracked'组合而成的。这个字符串通常用于表示跟踪的批次数。
if num_batches_tracked_key in state_dict:：检查状态字典中是否存在num_batches_tracked_key对应的键，即批次数是否在状态字典中。
del state_dict[num_batches_tracked_key]：如果存在num_batches_tracked_key，则从状态字典中删除这个键，因为在加载时不需要使用这个键的信息。
super(FrozenBatchNorm2d, self)._load_from_state_dict(...)：调用父类torch.nn.Module的_load_from_state_dict方法，实际上是从状态字典中加载模型的参数。这个调用确保了其他模型参数的正确加载。

总之，这个方法用于加载冻结的批量归一化层的状态字典，同时确保不加载用于跟踪批次数的信息，因为在推断过程中通常不需要这个信息。

3、forward()

    def forward(self, x):
        # move reshapes to the beginning
        # to make it fuser-friendly
        w = self.weight.reshape(1, -1, 1, 1)
        b = self.bias.reshape(1, -1, 1, 1)
        rv = self.running_var.reshape(1, -1, 1, 1)
        rm = self.running_mean.reshape(1, -1, 1, 1)
        eps = 1e-5
        scale = w * (rv + eps).rsqrt()
        bias = b - rm * scale
        return x * scale + bias

这是FrozenBatchNorm2d类的前向传播方法，用于在输入张量x上应用冻结的批量归一化。以下是代码的详细解释：

w = self.weight.reshape(1, -1, 1, 1)：首先，将归一化的权重张量self.weight重塑为形状为(1, C, 1, 1)的张量，其中C是通道数。这是为了与输入张量相乘。
b = self.bias.reshape(1, -1, 1, 1)：类似地，将偏置张量self.bias重塑为形状为(1, C, 1, 1)的张量。
rv = self.running_var.reshape(1, -1, 1, 1)：将运行方差张量self.running_var也重塑为形状为(1, C, 1, 1)的张量。
rm = self.running_mean.reshape(1, -1, 1, 1)：将运行均值张量self.running_mean同样重塑为形状为(1, C, 1, 1)的张量。
eps = 1e-5：定义了一个小的常数eps，用于避免除零错误。
scale = w * (rv + eps).rsqrt()：计算归一化的缩放因子。首先，对运行方差张量加上eps，然后计算其平方根的倒数。最后，将重塑后的权重张量w与这个倒数相乘，得到最终的缩放因子。
bias = b - rm * scale：计算偏置项，将重塑后的偏置张量b减去重塑后的运行均值张量rm与缩放因子scale的乘积。
return x * scale + bias：最后，将输入张量x与缩放因子scale相乘，并加上偏置项bias，以应用冻结的批量归一化。返回结果张量。

这个前向传播方法主要用于将输入数据归一化，然后通过可学习的缩放和偏置来调整数据，以便在网络中传播。这有助于提高网络的训练和推断性能。

三、BackboneBase()类

class BackboneBase(nn.Module):

    def __init__(self, backbone: nn.Module, train_backbone: bool, num_channels: int, return_interm_layers: bool):
        super().__init__()
        for name, parameter in backbone.named_parameters():
            if not train_backbone or 'layer2' not in name and 'layer3' not in name and 'layer4' not in name:
                parameter.requires_grad_(False)
        if return_interm_layers:
            return_layers = {"layer1": "0", "layer2": "1", "layer3": "2", "layer4": "3"}
        else:
            return_layers = {'layer4': "0"}
        self.body = IntermediateLayerGetter(backbone, return_layers=return_layers)
        self.num_channels = num_channels

    def forward(self, tensor_list: NestedTensor):
        xs = self.body(tensor_list.tensors)
        out: Dict[str, NestedTensor] = {}
        for name, x in xs.items():
            m = tensor_list.mask
            assert m is not None
            mask = F.interpolate(m[None].float(), size=x.shape[-2:]).to(torch.bool)[0]
            out[name] = NestedTensor(x, mask)
        return out

这是BackboneBase类的构造函数和前向传播方法，用于处理骨干网络（backbone）的特征提取。

1、init()

class BackboneBase(nn.Module):

    def __init__(self, backbone: nn.Module, train_backbone: bool, num_channels: int, return_interm_layers: bool):
        super().__init__()
        for name, parameter in backbone.named_parameters():
            if not train_backbone or 'layer2' not in name and 'layer3' not in name and 'layer4' not in name:
                parameter.requires_grad_(False)
        if return_interm_layers:
            return_layers = {"layer1": "0", "layer2": "1", "layer3": "2", "layer4": "3"}
        else:
            return_layers = {'layer4': "0"}
        self.body = IntermediateLayerGetter(backbone, return_layers=return_layers)
        self.num_channels = num_channels

构造函数 (__init__ 方法)：

def __init__(self, backbone: nn.Module, train_backbone: bool, num_channels: int, return_interm_layers: bool):：初始化函数接收四个参数：
- backbone：骨干网络的模块，通常是一个预训练的卷积神经网络，如ResNet。
- train_backbone：一个布尔值，指示是否训练整个骨干网络，如果为False，则只训练骨干网络的一部分。
- num_channels：表示骨干网络输出的特征通道数。
- return_interm_layers：一个布尔值，指示是否返回骨干网络中的中间层特征。
super().__init__()：调用父类nn.Module的构造函数进行初始化。
for name, parameter in backbone.named_parameters():：遍历骨干网络的命名参数。
if not train_backbone or 'layer2' not in name and 'layer3' not in name and 'layer4' not in name:：检查是否需要冻结骨干网络的参数。如果train_backbone为False，或者参数名称中不包含'layer2'、'layer3'和'layer4'，则将参数的梯度要求设为False，即不进行参数更新。这通常用于在微调模型时冻结骨干网络的一部分。
return_layers：根据return_interm_layers的值选择要返回的骨干网络的层。如果return_interm_layers为True，则返回"layer1"、"layer2"、"layer3"和"layer4"。如果为False，则仅返回"layer4"。
self.body = IntermediateLayerGetter(backbone, return_layers=return_layers)：创建IntermediateLayerGetter对象，该对象从骨干网络中获取指定的中间层。
self.num_channels = num_channels：将输出特征通道数存储在num_channels属性中，以供后续使用。

2、forward()

    def forward(self, tensor_list: NestedTensor):
        xs = self.body(tensor_list.tensors) #xs为最后一层特征图的信息
        out: Dict[str, NestedTensor] = {}
        for name, x in xs.items():
            m = tensor_list.mask
            assert m is not None
            mask = F.interpolate(m[None].float(), size=x.shape[-2:]).to(torch.bool)[0]
            out[name] = NestedTensor(x, mask)
        return out

前向传播方法 (forward 方法)：

def forward(self, tensor_list: NestedTensor):：前向传播方法接收一个NestedTensor对象作为输入，该对象包含输入张量和相应的掩码。
xs = self.body(tensor_list.tensors)：通过IntermediateLayerGetter对象将输入张量传递给骨干网络，以获取指定层的特征。xs是一个字典，包含不同层的特征。
out: Dict[str, NestedTensor] = {}：创建一个空字典out，用于存储输出特征。
for name, x in xs.items():：遍历xs字典中的每个层和相应的特征。
m = tensor_list.mask：从输入tensor_list中提取掩码信息。
assert m is not None：断言掩码不为空，确保掩码存在。
mask = F.interpolate(m[None].float(), size=x.shape[-2:]).to(torch.bool)[0]：将掩码插值到当前层的大小，以与特征对齐，并将结果转换为布尔类型。这确保了特征与输入具有相同的形状。
out[name] = NestedTensor(x, mask)：将特征张量x和相应的掩码mask存储在out字典中，作为输出的NestedTensor对象。
return out：返回包含不同层特征的字典out，其中键是层名称，值是NestedTensor对象，用于进一步的处理。

四、`Backbone`类的构造函数

class Backbone(BackboneBase):
    """ResNet backbone with frozen BatchNorm."""
    def __init__(self, name: str,
                 train_backbone: bool,
                 return_interm_layers: bool,
                 dilation: bool):
        #getattr返回torchvision.models的对象属性值
        backbone = getattr(torchvision.models, name)(
            replace_stride_with_dilation=[False, False, dilation],
            pretrained=is_main_process(), norm_layer=FrozenBatchNorm2d)
        num_channels = 512 if name in ('resnet18', 'resnet34') else 2048
        super().__init__(backbone, train_backbone, num_channels, return_interm_layers)

这是Backbone类的构造函数，它继承自BackboneBase类，并用于创建带有冻结BatchNorm的ResNet骨干网络。以下是代码的详细解释：

class Backbone(BackboneBase):：定义了名为Backbone的类，它继承自BackboneBase类，即Backbone是BackboneBase的子类。
def __init__(self, name: str, train_backbone: bool, return_interm_layers: bool, dilation: bool):：初始化函数接收四个参数：
- name：指定要使用的ResNet模型的名称，如'resnet18'或'resnet50'。
- train_backbone：一个布尔值，指示是否训练整个骨干网络，如果为False，则只训练骨干网络的一部分。
- return_interm_layers：一个布尔值，指示是否返回骨干网络中的中间层特征。
- dilation：一个布尔值，指示是否在ResNet的最后两个卷积层中使用空洞卷积（dilated convolution）。
backbone = getattr(torchvision.models, name)(...)：根据指定的ResNet模型名称name，使用getattr函数从torchvision.models模块中动态获取相应的ResNet模型。在创建ResNet模型时，传递了以下参数：
- replace_stride_with_dilation=[False, False, dilation]：这个参数控制是否替换ResNet的部分步幅（stride）为空洞卷积。dilation参数的值决定是否启用空洞卷积。
- pretrained=is_main_process()：这个参数控制是否使用预训练的权重。is_main_process()函数返回一个布尔值，用于确定当前进程是否为主进程。只有主进程加载预训练权重。
- norm_layer=FrozenBatchNorm2d：指定使用FrozenBatchNorm2d作为BatchNorm的层。
num_channels = 512 if name in ('resnet18', 'resnet34') else 2048：根据ResNet模型的名称name，确定输出特征的通道数num_channels。对于'resnet18'和'resnet34'，通道数为512；对于其他模型，通道数为2048。
super().__init__(backbone, train_backbone, num_channels, return_interm_layers)：调用父类BackboneBase的构造函数进行初始化，传递相应的参数。

总之，Backbone类的主要作用是创建指定的ResNet骨干网络，并根据参数配置是否冻结部分网络层，以及是否返回中间层特征。这个骨干网络将用于在目标检测模型中提取图像特征。

五、`Joiner`的PyTorch模型类

class Joiner(nn.Sequential):
    def __init__(self, backbone, position_embedding):
        super().__init__(backbone, position_embedding)

    def forward(self, tensor_list: NestedTensor):
        xs = self[0](tensor_list)
        out: List[NestedTensor] = []
        pos = []
        for name, x in xs.items():
            out.append(x)
            # position encoding
            pos.append(self[1](x).to(x.tensors.dtype))

        return out, pos

这是一个名为Joiner的PyTorch模型类，用于将骨干网络的输出特征与位置编码相结合。以下是代码的详细解释：

class Joiner(nn.Sequential):：定义了名为Joiner的类，它继承自nn.Sequential，这意味着Joiner类将包含一系列子模块，并且可以按顺序执行它们。
def __init__(self, backbone, position_embedding):：初始化函数接收两个参数：
- backbone：骨干网络模块，通常是一个PyTorch模型，它用于提取输入图像的特征。
- position_embedding：位置编码模块，通常是一个PyTorch模型，它用于为特征图添加位置信息。
super().__init__(backbone, position_embedding)：调用父类nn.Sequential的构造函数，将backbone和position_embedding添加到Joiner中，这样它们可以按顺序执行。
def forward(self, tensor_list: NestedTensor):：前向传播函数，接收一个NestedTensor对象作为输入，表示包含特征图和掩码的张量。
xs = self[0](tensor_list)：使用self[0]来执行第一个子模块，即backbone，将输入tensor_list传递给骨干网络。这将提取图像特征。
out: List[NestedTensor] = []：创建一个空列表out，用于存储特征图。
pos = []：创建一个空列表pos，用于存储位置编码信息。
for name, x in xs.items():：遍历骨干网络的输出字典，其中name是特征图的名称，x是对应的特征图。
out.append(x)：将特征图x添加到out列表中，以便后续的处理。
pos.append(self[1](x).to(x.tensors.dtype))：将特征图x传递给第二个子模块，即position_embedding，以生成位置编码。然后，将位置编码添加到pos列表中，并确保位置编码的数据类型与特征图的数据类型相匹配。
return out, pos：返回特征图列表out和位置编码列表pos。这些信息将在目标检测模型中用于进一步处理和预测目标框的位置和类别。

六、build_backbone()

def build_backbone(args):
    #两部分：backbone和位置编码，合在一起后传入至transformer中
    #对backbone输出的特征图进行位置编码，用于后续transformer部分
    position_embedding = build_position_encoding(args) #点进去看一下位置编码如何实现
    #是否训练backbone(即是否采用预训练backbone)
    train_backbone = args.lr_backbone > 0
    return_interm_layers = args.masks
    backbone = Backbone(args.backbone, train_backbone, return_interm_layers, args.dilation)
    #将backbone输出与position_embedding连在一起
    model = Joiner(backbone, position_embedding)
    model.num_channels = backbone.num_channels
    return model

这段代码是用于构建目标检测模型的骨干网络（backbone）和位置编码（position_embedding）的函数。下面是每行代码的详细解释：

position_embedding = build_position_encoding(args)：调用build_position_encoding函数，用于构建位置编码模块。位置编码是用于为特征图添加位置信息的关键组件。该函数返回一个位置编码模块，后续会与骨干网络连接在一起。
train_backbone = args.lr_backbone > 0：根据命令行参数lr_backbone的值，判断是否训练骨干网络。如果lr_backbone大于零，则表示需要对骨干网络进行微调，将train_backbone设置为True；否则，将其设置为False，表示不训练骨干网络。
return_interm_layers = args.masks：根据命令行参数masks的值，判断是否需要返回骨干网络的中间层特征。如果masks为True，则表示需要返回中间层特征，将return_interm_layers设置为True；否则，将其设置为False，表示只返回最后一层特征。
backbone = Backbone(args.backbone, train_backbone, return_interm_layers, args.dilation)：创建骨干网络模块，其中包括以下参数：
- args.backbone：命令行参数，指定要使用的骨干网络的名称（例如，ResNet）。
- train_backbone：上面判断的是否训练骨干网络的布尔值。
- return_interm_layers：上面判断的是否返回中间层特征的布尔值。
- args.dilation：命令行参数，指定是否采用空洞卷积（dilation）。
model = Joiner(backbone, position_embedding)：创建一个Joiner模块，将骨干网络模块（backbone）和位置编码模块（position_embedding）连接在一起。这个模块用于将特征图和位置编码整合在一起，以供后续的目标检测模型使用。
model.num_channels = backbone.num_channels：将model的num_channels属性设置为骨干网络的通道数量。这个属性用于指示特征图的通道数量。
return model：返回构建好的骨干网络模型，该模型包括了骨干网络和位置编码，可供目标检测模型使用。

#点进去看一下位置编码如何实现，转至文件models/position_encoding.py中。

别致的SmallSix