【扒代码】loca.py

最新推荐文章于 2024-09-03 17:37:40 发布

dearRongerr

最新推荐文章于 2024-09-03 17:37:40 发布

阅读量679

点赞数 17

分类专栏：扒代码文章标签： python 深度学习

本文链接：https://blog.csdn.net/2301_77549977/article/details/140994494

版权

扒代码专栏收录该内容

52 篇文章 1 订阅

订阅专栏

from .backbone import Backbone
from .transformer import TransformerEncoder
from .ope import OPEModule
from .positional_encoding import PositionalEncodingsFixed
from .regression_head import DensityMapRegressor

import torch
from torch import nn
from torch.nn import functional as F
# 忽略警告信息的代码
import warnings
warnings.filterwarnings("ignore")

import torch
from torch import nn

class LOCA(nn.Module):
    def __init__(
        self,
        # 模型配置参数
        image_size: int,                # 输入图像的尺寸
        num_encoder_layers: int,        # 编码器层数
        num_ope_iterative_steps: int,    # OPE迭代步骤数
        num_objects: int,              # 对象数量
        emb_dim: int,                  # 嵌入维度
        num_heads: int,                # 注意力机制中的头数
        kernel_dim: int,              # 卷积核尺寸
        backbone_name: str,            # 主干网络名称，如 'resnet50'
        swav_backbone: bool,           # 是否使用SWAV预训练的主干网络
        train_backbone: bool,          # 是否训练主干网络
        reduction: int,                # 特征融合时的降维因子
        dropout: float,                # dropout比率
        layer_norm_eps: float,        # 层归一化中的epsilon值
        mlp_factor: int,              # MLP的扩展因子
        norm_first: bool,              # 是否先进行归一化
        activation: nn.Module,         # 激活函数模块
        norm: bool,                    # 是否使用归一化
        zero_shot: bool,               # 是否是零样本学习场景
    ):
        super(LOCA, self).__init__()  # 调用基类的初始化方法

        # 初始化模型参数
        self.emb_dim = emb_dim
        self.num_objects = num_objects
        self.reduction = reduction
        self.kernel_dim = kernel_dim
        self.image_size = image_size
        self.zero_shot = zero_shot
        self.num_heads = num_heads
        self.num_encoder_layers = num_encoder_layers

        # 创建主干网络，可能是一个预训练的ResNet模型
        self.backbone = Backbone(
            backbone_name, pretrained=True, dilation=False, reduction=reduction,
            swav=swav_backbone, requires_grad=train_backbone
        )
        # 将主干网络的输出通道数映射到嵌入维度
        self.input_proj = nn.Conv2d(
            self.backbone.num_channels, emb_dim, kernel_size=1
        )

        # 如果编码器层数大于0，创建Transformer编码器
        if num_encoder_layers > 0:
            self.encoder = TransformerEncoder(
                num_encoder_layers, emb_dim, num_heads, dropout, layer_norm_eps,
                mlp_factor, norm_first, activation, norm
            )

        # 创建对象原型提取模块（OPEModule）
        self.ope = OPEModule(
            num_ope_iterative_steps, emb_dim, kernel_dim, num_objects, num_heads,
            reduction, layer_norm_eps, mlp_factor, norm_first, activation, norm, zero_shot
        )

        # 创建密度图回归头（用于计数）
        self.regression_head = DensityMapRegressor(emb_dim, reduction)
        # 创建辅助回归头列表，用于OPE迭代过程中的辅助任务
        self.aux_heads = nn.ModuleList([
            DensityMapRegressor(emb_dim, reduction)
            for _ in range(num_ope_iterative_steps - 1)
        ])

        # 创建位置编码模块
        self.pos_emb = PositionalEncodingsFixed(emb_dim)

功能解释：

LOCA 类是一个用于低样本对象计数的神经网络模型，它结合了多种组件和技术，如主干网络、Transformer编码器、对象原型提取（OPEModule）、密度图回归头等。
通过传入的参数，可以配置模型的不同方面，如嵌入维度、注意力头数、卷积核尺寸、主干网络类型、是否使用SWAV预训练权重、是否训练主干网络等。
Backbone 组件负责提取输入图像的特征，可能是一个预训练的ResNet模型，其参数根据配置进行初始化。
input_proj 是一个一维卷积层，用于将主干网络的输出映射到嵌入维度。
TransformerEncoder 组件实现了Transformer架构的编码器部分，用于进一步处理特征。
OPEModule 是对象原型提取模块，负责迭代地提取或适应对象原型。
DensityMapRegressor 组件用于预测密度图，从而实现对象计数。
PositionalEncodingsFixed 模块用于生成位置编码，提供序列或网格中每个元素的位置信息。

整体而言，LOCA 类实现了一个复杂的神经网络模型，用于在样本数量有限的情况下进行对象计数，它结合了多种先进的深度学习技术，以提高计数的准确性和鲁棒性。

# MLP的扩展因子扩展因子是啥

在机器学习和深度学习中，特别是在构建多层感知机（MLP）时，"扩展因子"（也称为宽度因子或扩展比例）是一个用于控制网络容量和大小的超参数。它定义了网络中层的输出维度与输入维度之间的比例关系。

具体来说，如果你有一个输入维度为 dindin 的层，并且你想应用一个扩展因子 mm 来增加其输出维度 doutdout，那么可以按照以下方式计算：

$d_{out} = m \times d_{in}$

这里的 m 就是扩展因子。使用扩展因子可以很容易地调整网络的宽度，而不需要手动指定每一层的输出维度。

例如，如果输入维度是 64，你选择的扩展因子是 2，那么输出维度将是 128。这在设计网络时非常有用，因为它允许模型设计者通过一个简单的乘法因子来放大或缩小网络的宽度。

在某些情况下，扩展因子还可以用于控制非线性激活函数的增长速率，或者在构建具有多个隐藏层的网络时，逐步增加或减少每一层的输出维度。

在提供的代码上下文中，mlp_factor 用于确定多层感知机中间层的尺寸。如果 mlp_factor 是 2，而输入特征维度是 emb_dim，那么中间层的维度将是 2 * emb_dim。这有助于模型学习更复杂的函数映射，但同时也需要权衡模型的容量和避免过拟合的风险。

forward函数

def forward(self, x, bboxes):
    # 确定对象的数量，如果不是零样本学习场景，则根据bboxes的数量确定
    num_objects = bboxes.size(1) if not self.zero_shot else self.num_objects

    # 通过主干网络提取特征
    backbone_features = self.backbone(x)

    # 准备编码器的输入
    src = self.input_proj(backbone_features)
    # 获取特征的尺寸
    bs, c, h, w = src.size()

    # 生成位置编码并调整其形状以匹配编码器的输入
    pos_emb = self.pos_emb(bs, h, w, src.device).flatten(2).permute(2, 0, 1)
    # 调整src的形状
    src = src.flatten(2).permute(2, 0, 1)

    # 通过编码器处理特征
    if self.num_encoder_layers > 0:
        image_features = self.encoder(src, pos_emb, src_key_padding_mask=None, src_mask=None)
    else:
        image_features = src

    # 准备OPE（对象原型提取）模块的输入
    f_e = image_features.permute(1, 2, 0).reshape(-1, self.emb_dim, h, w)

    # 调用OPE模块生成所有原型
    all_prototypes = self.ope(f_e, pos_emb, bboxes)

    # 初始化输出列表
    outputs = list()

    # 遍历所有原型
    for i in range(all_prototypes.size(0)):
        # 处理每个原型
        prototypes = all_prototypes[i, ...].permute(1, 0, 2).reshape(
            bs, num_objects, self.kernel_dim, self.kernel_dim, -1
        ).permute(0, 1, 4, 2, 3).flatten(0, 2)[:, None, ...]

        # 使用原型和查询特征生成响应图
        response_maps = F.conv2d(
            torch.cat([f_e for _ in range(num_objects)], dim=1).flatten(0, 1).unsqueeze(0),
            prototypes,
            bias=None,
            padding=self.kernel_dim // 2,
            groups=prototypes.size(0)
        ).view(
            bs, num_objects, self.emb_dim, h, w
        ).max(dim=1)[0]

        # 通过回归头处理响应图
        if i == all_prototypes.size(0) - 1:
            predicted_dmaps = self.regression_head(response_maps)
        else:
            predicted_dmaps = self.aux_heads[i](response_maps)
        # 将预测的密度图添加到输出列表
        outputs.append(predicted_dmaps)

    # 返回最终的预测密度图和中间辅助输出
    return outputs[-1], outputs[:-1]

功能解释：

forward 方法接收输入图像 x 和边界框 bboxes，然后通过一系列的处理步骤生成预测的密度图。
backbone_features 是主干网络提取的特征。
src 是经过投影层 input_proj 处理的编码器输入。
pos_emb 是位置编码，它提供了关于特征位置的信息。
image_features 是编码器的输出特征。
all_prototypes 是 OPE 模块生成的对象原型集合。
通过遍历 all_prototypes 中的每个原型，使用 2D 卷积 F.conv2d 生成响应图 response_maps。
predicted_dmaps 是通过回归头处理响应图后得到的预测密度图。
方法返回最终的预测密度图和中间辅助输出，这些输出可以用于模型训练过程中的损失计算或分析。

关于 `all_prototypes` 的问题：

all_prototypes 是由 OPE 模块生成的对象原型集合，它包含了用于与查询图像特征进行匹配的多个原型表示。每个原型都是通过迭代适应过程生成的，用于捕捉图像中目标对象的形状和外观特征。在循环中，每个原型被用来生成响应图，这些响应图随后被送入回归头以产生预测的密度图。

MyDoubt

NOTE：所以我到底应该改改哪里呢？

没看到怎么从数据开始处理特征的

充分考虑考虑位置特征和外观特征

主要还是不太明白网络架构和数据流动过程

`build_model`

名为 build_model 的函数，它用于根据提供的参数 args 构建并返回一个 LOCA 模型实例

def build_model(args):
    # 确保提供的参数中，backbone的名称是支持的ResNet版本之一
    assert args.backbone in ['resnet18', 'resnet50', 'resnet101'], "Unsupported backbone name"
    # 确保提供的参数中，reduction的值是支持的降维因子之一
    assert args.reduction in [4, 8, 16], "Unsupported reduction factor"

    # 根据参数构建并返回一个LOCA模型实例
    return LOCA(
        # 以下参数均从args对象中获取，并传递给LOCA模型的构造函数
        image_size=args.image_size,  # 输入图像的尺寸
        num_encoder_layers=args.num_enc_layers,  # 编码器层数
        num_ope_iterative_steps=args.num_ope_iterative_steps,  # OPE迭代步骤数
        num_objects=args.num_objects,  # 对象数量
        zero_shot=args.zero_shot,  # 是否是零样本学习场景
        emb_dim=args.emb_dim,  # 嵌入维度
        num_heads=args.num_heads,  # 注意力机制中的头数
        kernel_dim=args.kernel_dim,  # 卷积核尺寸
        backbone_name=args.backbone,  # 主干网络名称
        swav_backbone=args.swav_backbone,  # 是否使用SWAV预训练的主干网络
        # 如果提供了非零的主干网络学习率，则训练主干网络
        train_backbone=args.backbone_lr > 0,
        reduction=args.reduction,  # 特征融合时的降维因子
        dropout=args.dropout,  # dropout比率
        layer_norm_eps=1e-5,  # 层归一化中的epsilon值
        mlp_factor=8,  # MLP的扩展因子
        norm_first=args.pre_norm,  # 是否先进行归一化
        activation=nn.GELU,  # 激活函数使用GELU
        norm=True,  # 是否使用归一化
    )

功能解释：

build_model 函数首先通过断言（assert）语句确保 args 中的 backbone 和 reduction 参数是有效的。
然后，函数创建并返回一个 LOCA 类的实例，传递了一系列参数来配置模型的行为，包括网络结构、训练选项、正则化策略等。
args 对象被假定为包含所有必要的配置参数，这些参数被用来初始化 LOCA 模型。
LOCA 模型是用于低样本对象计数任务的神经网络模型，它结合了多种组件和技术，如主干网络、Transformer编码器、对象原型提取模块等。

整体而言，build_model 函数提供了一个简洁的方式来创建定制化的 LOCA 模型实例，使得模型可以根据不同的实验设置或应用需求进行调整。

from .backbone import Backbone
from .transformer import TransformerEncoder
from .ope import OPEModule
from .positional_encoding import PositionalEncodingsFixed
from .regression_head import DensityMapRegressor

import torch
from torch import nn
from torch.nn import functional as F
# 忽略警告信息的代码
import warnings
warnings.filterwarnings("ignore")

class LOCA(nn.Module):

    '''
    实现低样本对象计数
    （Low-shot Object Counting with iterative prototype Adaptation）模型
    '''

    def __init__(
        self,
        image_size: int,# 输入图像的尺寸
        num_encoder_layers: int,# 编码器层数
        num_ope_iterative_steps: int,#OPE迭代步骤数
        num_objects: int,# 对象数量
        emb_dim: int,# 嵌入维度
        num_heads: int,# 注意力机制中的头数
        kernel_dim: int,# 卷积核尺寸
        backbone_name: str,# 主干网络名称，如 'resnet50'
        swav_backbone: bool,# 是否使用SWAV预训练的主干网络
        train_backbone: bool,# 是否训练主干网络
        reduction: int,# 特征融合时的降维因子
        dropout: float,# dropout比率
        layer_norm_eps: float,# 层归一化中的epsilon值
        mlp_factor: int,# MLP的扩展因子
        norm_first: bool,# 是否先进行归一化 
        activation: nn.Module,# 激活函数模块
        norm: bool,# 是否使用归一化
        zero_shot: bool,# 是否是零样本学习场景
    ):

        super(LOCA, self).__init__()# 调用基类的初始化方法

        # 初始化模型参数
        self.emb_dim = emb_dim
        self.num_objects = num_objects
        self.reduction = reduction
        self.kernel_dim = kernel_dim
        self.image_size = image_size
        self.zero_shot = zero_shot
        self.num_heads = num_heads
        self.num_encoder_layers = num_encoder_layers

        # 创建主干网络，可能是一个预训练的ResNet模型
        self.backbone = Backbone(
            backbone_name, pretrained=True, dilation=False, reduction=reduction,
            swav=swav_backbone, requires_grad=train_backbone
        )

        # 将主干网络的输出通道数映射到嵌入维度
        self.input_proj = nn.Conv2d(
            self.backbone.num_channels, emb_dim, kernel_size=1
        )

        # 如果编码器层数大于0，创建Transformer编码器
        if num_encoder_layers > 0:
            self.encoder = TransformerEncoder(
                num_encoder_layers, emb_dim, num_heads, dropout, layer_norm_eps,
                mlp_factor, norm_first, activation, norm
            )


        # 创建对象原型提取模块（OPEModule）
        self.ope = OPEModule(
            num_ope_iterative_steps, emb_dim, kernel_dim, num_objects, num_heads,
            reduction, layer_norm_eps, mlp_factor, norm_first, activation, norm, zero_shot
        )

        # 创建密度图回归头（用于计数）
        self.regression_head = DensityMapRegressor(emb_dim, reduction)
        # 创建辅助回归头列表，用于OPE迭代过程中的辅助任务
        self.aux_heads = nn.ModuleList([
            DensityMapRegressor(emb_dim, reduction)
            for _ in range(num_ope_iterative_steps - 1)
        ])

         # 创建位置编码模块
        self.pos_emb = PositionalEncodingsFixed(emb_dim)

    def forward(self, x, bboxes):
        # 确定对象的数量，如果不是零样本学习场景，则根据bboxes的数量确定
        num_objects = bboxes.size(1) if not self.zero_shot else self.num_objects
        # backbone
        # 通过主干网络提取特征
        backbone_features = self.backbone(x)
        # prepare the encoder input
        # 准备编码器的输入
        src = self.input_proj(backbone_features)
        # 获取特征的尺寸
        bs, c, h, w = src.size()
        # 生成位置编码并调整其形状以匹配编码器的输入
        pos_emb = self.pos_emb(bs, h, w, src.device).flatten(2).permute(2, 0, 1)
        # 调整src的形状
        src = src.flatten(2).permute(2, 0, 1)

        # push through the encoder
        # 通过编码器处理特征
        if self.num_encoder_layers > 0:
            image_features = self.encoder(src, pos_emb, src_key_padding_mask=None, src_mask=None)
        else:
            image_features = src

        # prepare OPE input
        # 准备OPE（对象原型提取）模块的输入
        f_e = image_features.permute(1, 2, 0).reshape(-1, self.emb_dim, h, w)

        # 调用OPE模块生成所有原型
        # 问题：原型到底是啥
        all_prototypes = self.ope(f_e, pos_emb, bboxes)

        # 初始化输出列表
        outputs = list()

        # 遍历所有原型
        # 问题： 这里的all_prototypes是什么？
        for i in range(all_prototypes.size(0)):
            # 处理每个原型
            prototypes = all_prototypes[i, ...].permute(1, 0, 2).reshape(
                bs, num_objects, self.kernel_dim, self.kernel_dim, -1
            ).permute(0, 1, 4, 2, 3).flatten(0, 2)[:, None, ...]

            # 使用原型和查询特征生成响应图
            # 原型：prototypes
            # 查询特征：[f_e for _ in range(num_objects)]
            response_maps = F.conv2d(
                torch.cat([f_e for _ in range(num_objects)], dim=1).flatten(0, 1).unsqueeze(0),
                prototypes,
                bias=None,
                padding=self.kernel_dim // 2,
                groups=prototypes.size(0)
            ).view(
                bs, num_objects, self.emb_dim, h, w
            ).max(dim=1)[0]

            # send through regression heads
            # 通过回归头处理响应图
            if i == all_prototypes.size(0) - 1:
                predicted_dmaps = self.regression_head(response_maps)
            else:
                predicted_dmaps = self.aux_heads[i](response_maps)
            # 将预测的密度图添加到输出列表
            outputs.append(predicted_dmaps)

        # 返回最终的预测密度图和中间辅助输出
        return outputs[-1], outputs[:-1]


def build_model(args):

    # 通过断言（assert）语句确保 args 中的 backbone 和 reduction 参数是有效的
    assert args.backbone in ['resnet18', 'resnet50', 'resnet101']
    assert args.reduction in [4, 8, 16]

    # 返回一个LOCA模型实例
    return LOCA(
        image_size=args.image_size,# 输入图像的尺寸
        num_encoder_layers=args.num_enc_layers,# 编码器层数
        num_ope_iterative_steps=args.num_ope_iterative_steps,# OPE迭代步骤数
        num_objects=args.num_objects,# 对象数量
        zero_shot=args.zero_shot,# 是否是零样本学习场景
        emb_dim=args.emb_dim,# 嵌入维度
        num_heads=args.num_heads,# 注意力机制中的头数
        kernel_dim=args.kernel_dim,# 卷积核尺寸
        backbone_name=args.backbone,# 主干网络名称
        swav_backbone=args.swav_backbone,# 是否使用SWAV预训练的主干网络
        train_backbone=args.backbone_lr > 0,# 如果提供了非零的主干网络学习率，则训练主干网络
        reduction=args.reduction,# 特征融合时的降维因子
        dropout=args.dropout,# dropout比率
        layer_norm_eps=1e-5,# 层归一化中的epsilon值
        mlp_factor=8,# MLP的扩展因子
        norm_first=args.pre_norm,# 是否先进行归一化
        activation=nn.GELU,# 激活函数使用GELU
        norm=True,# 是否使用归一化
    )