LoFTR源码详解+个人对LoFTR的细节之处的理解

摘要

本文主要是针对LoFTR架构，对其源码进行debug详细分析，从而更进一步加深对LoFTR架构的理解。

1. LoFTR整体架构分析

LoFTR论文详解请参见：LoFTR论文详解（特征匹配），LoFTR源码中对于LoFTR架构的整体处理过程代码在 src/loftr.py
接下来我们首先分析 loftr,py 中 LoFTR 类初始化。整体流程就是，首先通过 resnet+fpn 主干网络进行两幅图像的特征提取，得到粗粒度特征（尺寸大小为原输入图像的 1/8）和细粒度特征（尺寸大小为原输入图像的 1/2），将粗粒度特征进行位置编码后，作为粗粒度特征匹配中的 LoFTR 模块（自注意力模块+交叉注意力模块），这样得到的各个全局关联粗粒度特征，将两幅图像经过Transformer处理后的特征再进行粗粒度特征匹配模块，确定两幅图像所能匹配到的关键点个数以及关键点坐标（相对于原图像分辨率的坐标），但是经过粗粒度确定的关键点坐标只是在一个个小区域所确定的坐标，并不精确，因此需要细粒度特征处理模块进行细化粗粒度特征，然后将最初特征提取得到的两幅图像的细粒度特征进行裁剪，裁剪至 5*5 大小，根据粗粒度特征匹配中确定下来的大致的关键点坐标，得到细粒度特征匹配的大致关键点，接着沿批次数将只有这些关键点的两组细粒度特征进行拼接，然后将粗粒度匹配得到的两组张量也进行拼接，随后将拼接后的粗粒度特征与细粒度特征进行特征融合，接着将融合后的特征张量分割成两部分，分别作为细粒度特征匹配中的LoFTR模块的输入数据，最后得到一个二维的热图，然后计算归一化坐标，得到最终的两幅图像匹配的关键点坐标。

    def __init__(self, config):
        super().__init__()
        # LoFTR所需的配置参数
        self.config = config

        # 特征提取的主干网络
        self.backbone = build_backbone(config)
        # 位置编码模块
        self.pos_encoding = PositionEncodingSine(
            config['coarse']['d_model'],
            temp_bug_fix=config['coarse']['temp_bug_fix'])
        # 粗粒度特征处理
        self.loftr_coarse = LocalFeatureTransformer(config['coarse'])
        # 建立粗粒度特征匹配
        self.coarse_matching = CoarseMatching(config['match_coarse'])
        # 细粒度特征处理 + 粗粒度特征与细粒度特征融合
        self.fine_preprocess = FinePreprocess(config)
        # 细粒度部分的 LoFTR 模块
        self.loftr_fine = LocalFeatureTransformer(config["fine"])
        # 确定最终匹配的特征细化关键点坐标
        self.fine_matching = FineMatching()

LoFTR 类的前向传播过程

1. 首先使用data.update()方法来更新字典data其中bs是新添加到data字典的键，代表批量大小（batch size)。data[‘image0’].size(0)获取data字典中’image0’键对应张量的第一个维度的大小，即批量大小，并将其赋值给’bs’。‘hw0_i’ 和 ‘hw1_i’ 是新添加到 data 字典的键，用于存储两幅图像的高度和宽度信息。data[‘image0’].shape[2:] 和 data[‘image1’].shape[2:] 分别获取 data 字典中 ‘image0’ 和 ‘image1’ 键对应张量的第2维和第3维的大小，这通常对应于图像的高度和宽度，并将其分别赋值给 ‘hw0_i’ 和 ‘hw1_i’。

 data.update({
            'bs': data['image0'].size(0),
            'hw0_i': data['image0'].shape[2:], 'hw1_i': data['image1'].shape[2:]
        })

此时得到的 data 字典为：

2. 如果两幅图像的高度和宽度相同，执行条件块内的代码。将两幅图像沿第0维（批量大小维度）拼接，然后使用 self.backbone 方法处理拼接后的图像，提取特征。然后将特征张量feats_c和feats_f按批量大小分割成两部分，分别对应两幅图像的特征。即粗粒度特征和细粒度特征，如果两幅图像的高度和宽度不相同，执行条件块内else的代码，分别对两幅图像调用 self.backbone 方法提取特征。

        if data['hw0_i'] == data['hw1_i']:  # faster & better BN convergence
            feats_c, feats_f = self.backbone(torch.cat([data['image0'], data['image1']], dim=0))
            (feat_c0, feat_c1), (feat_f0, feat_f1) = feats_c.split(data['bs']), feats_f.split(data['bs'])
        else:  # handle different input shapes
            (feat_c0, feat_f0), (feat_c1, feat_f1) = self.backbone(data['image0']), self.backbone(data['image1'])

再次更新data字典，这次添加了特征张量的形状信息，其中hw0_i、hw1_i表示的是输入数据image0、image1的尺寸大小，hw0_c、hw1_c表示两幅图像的粗粒度特征，hw0_f、hw1_f表示的是两幅图像的细粒度特征。

 data.update({
            'hw0_c': feat_c0.shape[2:], 'hw1_c': feat_c1.shape[2:],
            'hw0_f': feat_f0.shape[2:], 'hw1_f': feat_f1.shape[2:]
        })

接着将粗粒度特征 feat_c0、feat_c1 进行位置编码

 feat_c0 = rearrange(self.pos_encoding(feat_c0), 'n c h w -> n (h w) c')
 feat_c1 = rearrange(self.pos_encoding(feat_c1), 'n c h w -> n (h w) c')

然后为image0和image1的粗粒度张量加入mask掩码遮罩，以保证网络架构的准确性，创建mask_c0、mask_c1来保存image0、image1粗粒度特征的掩码，如果mask_c0、mask_c1为空，调用 self.loftr_coarse 方法，传入特征张量feat_c0和feat_c1，以及对应的掩码mask_c0和mask_c1。

 mask_c0 = mask_c1 = None  # mask is useful in training
 if 'mask0' in data:
     mask_c0, mask_c1 = data['mask0'].flatten(-2), data['mask1'].flatten(-2)
 feat_c0, feat_c1 = self.loftr_coarse(feat_c0, feat_c1, mask_c0, mask_c1)

进行粗粒度特征匹配

 self.coarse_matching(feat_c0, feat_c1, data, mask_c0=mask_c0, mask_c1=mask_c1)

进行粗粒度特征细化，以及粗粒度特征和细粒度特征融合

feat_f0_unfold, feat_f1_unfold = self.fine_preprocess(feat_f0, feat_f1, feat_c0, feat_c1, data)
        if feat_f0_unfold.size(0) != 0:  # at least one coarse level predicted
            feat_f0_unfold, feat_f1_unfold = self.loftr_fine(feat_f0_unfold, feat_f1_unfold)

确定最终的关键点坐标

self.fine_matching(feat_f0_unfold, feat_f1_unfold, data)片

整体 LoFTR 前向传播代码：

 def forward(self, data):
        """ 
        Update:
            data (dict): {
                'image0': (torch.Tensor): (N, 1, H, W)
                'image1': (torch.Tensor): (N, 1, H, W)
                'mask0'(optional) : (torch.Tensor): (N, H, W) '0' indicates a padded position
                'mask1'(optional) : (torch.Tensor): (N, H, W)
            }
        """
        print(data)
        # 1. Local Feature CNN
        data.update({
            'bs': data['image0'].size(0),
            'hw0_i': data['image0'].shape[2:], 'hw1_i': data['image1'].shape[2:]
        })
        print(data)
        if data['hw0_i'] == data['hw1_i']:  # faster & better BN convergence
            feats_c, feats_f = self.backbone(torch.cat([data['image0'], data['image1']], dim=0))
            (feat_c0, feat_c1), (feat_f0, feat_f1) = feats_c.split(data['bs']), feats_f.split(data['bs'])
        else:  # handle different input shapes
            (feat_c0, feat_f0), (feat_c1, feat_f1) = self.backbone(data['image0']), self.backbone(data['image1'])

        data.update({
            'hw0_c': feat_c0.shape[2:], 'hw1_c': feat_c1.shape[2:],
            'hw0_f': feat_f0.shape[2:], 'hw1_f': feat_f1.shape[2:]
        })

        # 2. coarse-level loftr module
        # add featmap with positional encoding, then flatten it to sequence [N, HW, C]
        feat_c0 = rearrange(self.pos_encoding(feat_c0), 'n c h w -> n (h w) c')
        feat_c1 = rearrange(self.pos_encoding(feat_c1), 'n c h w -> n (h w) c')

        mask_c0 = mask_c1 = None  # mask is useful in training
        if 'mask0' in data:
            mask_c0, mask_c1 = data['mask0'].flatten(-2), data['mask1'].flatten(-2)
        feat_c0, feat_c1 = self.loftr_coarse(feat_c0, feat_c1, mask_c0, mask_c1)

        # 3. match coarse-level
        self.coarse_matching(feat_c0, feat_c1, data, mask_c0=mask_c0, mask_c1=mask_c1)

        # 4. fine-level refinement
        feat_f0_unfold, feat_f1_unfold = self.fine_preprocess(feat_f0, feat_f1, feat_c0, feat_c1, data)
        if feat_f0_unfold.size(0) != 0:  # at least one coarse level predicted
            feat_f0_unfold, feat_f1_unfold = self.loftr_fine(feat_f0_unfold, feat_f1_unfold)

        # 5. match fine-level
        self.fine_matching(feat_f0_unfold, feat_f1_unfold, data)

接下来，我将对各个模块进行详细 debug 分析。

2. 主干网络

原论文中只是简要说了LoFTR采用CNN进行特征提取，而在debug源码后，发现LoFTR采用resnet+fpn进行特征提取，整体特征提取的网络是将resnet与FPN整合到一起了，由于论文中描述的LoFTR的思路是首先以低特征分辨率（图像维度的1/8）在两组变换特征之间提取密集匹配。从这些密集匹配中选择具有高置信度的匹配，然后使用基于相关性的方法将其细化到子像素级。而源码中正是通过了一轮的卷积+三层的resnet BasicBlock下面是第一轮卷积的推导过程：
第一轮卷积：

第一轮只进行了卷积、批量归一化、激活函数，第二轮到第四轮则是通过resnet残差block来进行处理数据，代码中将其分为layer1、layer2、layer3，其中layer1中卷积的步长为1，这可以说明当数据从第一层输出后，layer1进行处理后，数据的宽高是不变的，而相对于原始输入的数据，经过layer1后，宽高仍然为原输入图像的二分之一，论文中说到在进行特征处理的时候进行粗粒度提取特征（原始图像的八分之一）、将原始图像大小二分之一作为细粒度特征，细粒度特征后续会用到，这里不过多赘述。图像数据通过layer1、layer2、layer3计算思路同上。

这里注意到每个layer中有两个BasicBlock残差块，源码中写到

在这里我手推计算下第二个layer2中数据通过第二个BasicBlock的尺度变化，因为layer2中二者的步长不同。

由此可见，经过第二个BasicBlock尺度不变，这么做的目的是为了更好的提取特征，减少网络过拟合，下图便是layer2网络结构

Layer3网络同layer2，这里注意到，layer2网络中多出了一个downsample下采样层，而该下采样层中的 1 * 1 卷积的步长为2、padding为0，通过卷积计算后，维度与 y 的维度相同，长宽不变，x 则是上一层输出的数据，其实这里也可以直接将输入的数据x与y进行相加，但是这里使用 1*1 卷积进行处理，好处就是在保证数据尺度不变的前提下增加了非线性、减少了参数，二次提取特征。然后与该层中通过卷积的数据进行融合，源码中BasicBlock写到

BasicBlock类的前向传播：

 def forward(self, x):
        y = x
        y = self.relu(self.bn1(self.conv1(y)))
        y = self.bn2(self.conv2(y))

        if self.downsample is not None:
            x = self.downsample(x)

        return self.relu(x+y)

以下便是resnet+fpn里，数据最先通过resnet后的形状和维度变化：

因此，同论文中所阐述的，将其分为粗粒度特征和细粒度特征，然后将x3（对应最初输入图像大小的八分之一）进行上采样，得到x3_out_2x，将x3_out_2x（上采样后的数据）与x2_out进行特征融合，得到x2_out，然后对x2_out进行上采样得到x2_out_2x，

见代码：

 # FPN
        x3_out = self.layer3_outconv(x3)

        x3_out_2x = F.interpolate(x3_out, scale_factor=2., mode='bilinear', align_corners=True)
        x2_out = self.layer2_outconv(x2)
        x2_out = self.layer2_outconv2(x2_out+x3_out_2x)

        x2_out_2x = F.interpolate(x2_out, scale_factor=2., mode='bilinear', align_corners=True)
        x1_out = self.layer1_outconv(x1)
        x1_out = self.layer1_outconv2(x1_out+x2_out_2x)
        print(x.shape)
        return [x3_out, x1_out]

其中F.interpolate函数对x3_out进行上采样，scale_factor=2.表示将特征图的每个维度放大2倍。mode='bilinear’指定了使用双线性插值作为上采样的方法。align_corners=True是一个参数，用于控制插值时角落像素的对齐方式。

x3_out、x2_out 是经过11卷积层处理后的特征，用于提取更高级的特征表示。x3上采样后得到x3_out_2x，将其与x2_out进行特征融合，然后将x2_out作为输入，进行上采样得到 x2_out_2x，由于原先的layer层只有三层，故 x1 在特征金字塔中属于顶层，无需上采样了，只需将x1进行11卷积后，丰富其特征，然后将其与x2上采样的数据进行特征融合。最后得到的x3_out、x1_out分别是粗粒度特征（原输入图像尺寸的1/8）、细粒度特征（原输入图像尺寸的1/2）。

3. 位置编码

首先创建一个形状为 (d_model, height, width) 的全零张量pe，用于存储位置编码。然后分别创建与 max_shape 相同形状的全1矩阵、全0矩阵，沿第0维（高度）累加求和，转换为浮点数，并增加一个维度，接着生成上面公式的除数项div_term。由于位置编码中对于奇数列用cos()函数计算，偶数列用sin()函数计算，而源码中是将原有的pe矩阵里的256分成4部分，分别进行计算

通过这种处理得到位置编码矩阵pe，将进行特征提取后的图像数据（粗粒度c、细粒度f）加上位置编码, 此时x0=（1，256，60，80）（image0的粗粒度特征）、x1=(1,256,60,80)（image1的粗粒度特征），然后将image0和image1的粗粒度张量的宽w和高h相乘，将原本的四维张量加上位置编码后转换为三维张量，作为Transformer架构的输入数据

loftr.py 前向传播中位置编码部分：

转换前的image0和image1的粗粒度张量转换后的image0和image1的粗粒度张量

4. 粗粒度特征中的Transformer

回到 loftr.py 的前向传播中，接下来为image0和image1的粗粒度张量加入mask掩码遮罩，以保证网络架构的准确性，创建mask_c0、mask_c1来保存image0、image1粗粒度特征的掩码，如果mask_c0、mask_c1为空，调用 self.loftr_coarse 方法，传入特征张量feat_c0和feat_c1，以及对应的掩码mask_c0和mask_c1。

对应的loftr_coarse方法如下，loftr_coarse方法便是LoFTR架构中的Transformer模块
根据整体代码的参数配置，参数配置文件在cvpr_ds_config.py中，LoFTR-coarse module config内的d_model为256，nhead为8，layer_names为[‘self’, ‘cross’] * 4，attention为linear。其中LocalFeatureTransformer类中的d_model表示维度大小，nhead表示这个粗粒度特征对应的transformer共有8层，layer_names表示的是自注意力还是交叉注意力，encoder_layer表示的是编码器层，self.layers是用于存储多个编码器模块。

根据 LoFTR 的参数设定，对于每副图像的编码器模块一共有8层，下面的代码则是 LocalFeatureTransformer类的整体代码，LocalFeatureTransformer主要是定义Transformer结构。

class LocalFeatureTransformer(nn.Module):
    """A Local Feature Transformer (LoFTR) module."""

    def __init__(self, config):
        super(LocalFeatureTransformer, self).__init__()

        self.config = config
        self.d_model = config['d_model']
        self.nhead = config['nhead']
        self.layer_names = config['layer_names']
        encoder_layer = LoFTREncoderLayer(config['d_model'], config['nhead'], config['attention'])
        self.layers = nn.ModuleList([copy.deepcopy(encoder_layer) for _ in range(len(self.layer_names))])
        self._reset_parameters()

    def _reset_parameters(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)
        """
        Args:
            feat0 (torch.Tensor): [N, L, C]
            feat1 (torch.Tensor): [N, S, C]
            mask0 (torch.Tensor): [N, L] (optional)
            mask1 (torch.Tensor): [N, S] (optional)
        """

    def forward(self, feat0, feat1, mask0=None, mask1=None):
        assert self.d_model == feat0.size(2)
        for layer, name in zip(self.layers, self.layer_names):
            if name == 'self':
                feat0 = layer(feat0, feat0, mask0, mask0)
                feat1 = layer(feat1, feat1, mask1, mask1)
            elif name == 'cross':
                feat0 = layer(feat0, feat1, mask0, mask1)
                feat1 = layer(feat1, feat0, mask1, mask0)
            else:
                raise KeyError

        return feat0, feat1

接着便是涉及编码器LoFTREncoderLayer部分的细节：

编码器由8个相同的modellist 组成。第一层是多头自注意力机制，第二层是位置全连接前馈网络，然后分别创建三个线性层，用于计算q、k、v向量，根据传入的 attention 参数类型，实例化一个线性注意力层 LinearAttention 或者全注意力层 FullAttention。接着创建一个线性层，用于将多头注意力的输出合并回 d_model 维度。

接下来便是前馈网络，创建一个顺序容器 Sequential，包含两个线性层和一个 ReLU 激活函数，用于构建前馈网络（MLP，多层感知机）。前馈网络的第一个线性层，输入和输出维度都是 d_model 的两倍。前馈网络的第二个线性层，将特征维度从 d_model 的两倍降维回 d_model。最后将输出进行层归一化。

通过层归一化每个样本的特征都被独立归一化，不受批次大小影响。可以减少梯度消失问题，加速网络训练。有助于稳定训练过程，提高模型的泛化能力。层归一化在深度学习中得到广泛应用，特别是在自然语言处理任务和Transformer模型中，有助于提高模型的效果和训练速度。

由于我们现在分析的是编码器的第一层，而自注意力模块和交叉注意力模块是交替处理数据的，第一层为自注意力，故输入的张量均为 feat0 或 feat1，然后为多头注意力机制准备数据，query、key、value原本的值shape为（1，4800，256），其中bs为1，图像数据序列长度为4800，维度为256。Q、k、v的 view 操作将查询张量重新排列为 [batch_size, sequence_length, nhead, dim] 的形状，得到的shape为（1，4800，8，32），sequence_length 被分割成 8 组，每组具有 32 个特征。这样，每个头可以独立地处理输入序列的一部分特征，通过将q、k、v张量分割成多个头，每个头可以学习到输入数据的不同方面，然后将这些信息合并起来，以获得更全面的表示。
有了q、k、v三个张量后，将其作为多头注意力层linear_attention的输入数据。

LinearAttention 类 用于实现一种线性复杂度的注意力机制，这在处理长序列时特别有用，因为它可以减少计算量。类中的 self.feature_map 函数用于在计算注意力分数之前对输入特征进行非线性变换，而 self.eps 参数用于在进行 softmax 操作或其他涉及除法的操作时提供数值稳定性。

函数体返回对输入张量 x 应用 ELU 激活函数的结果，并在此基础上加 1，ELU是非线性的，可以帮助神经网络学习复杂的数据表示
Linear_attention前向传播过程：

Linear_attention 整体思路

将Q和K的长度表示为N并且将它们的特征维数表示为D，则在Transformer中Q和K之间的点积引入了计算成本，该计算成本随着输入序列的长度而二次增长（$O(n^2)$）。在局部特征匹配的上下文中直接应用普通版本的Transformer是不切实际的。为了解决这一问题，作者在Transformer中使用 vanilla 注意层的有效变体。Linear Transformer 提出将Transformer的计算复杂度降低到 $O(n)$，Linear Transformer 通过用替代核函数 $sim(Q,K)=\phi (Q)\phi (K{)}^T$替换原始注意力层中使用的指数核，其中 φ（·）= elu（·）+1，利用矩阵乘积的结合性，可以先进行$\phi (K{)}^T$与$V$的乘法运算。因为D << N，所以计算成本降低到O（N）。

首先将上述我们q、k、v三个张量中q、k通过ELU激活函数进行非线性激活，由于在编码器中均不设置掩码mask，编码器的输出会被传递到解码器，而不会进行掩码操作。

接下来将k、v张量进行矩阵乘法，其中 K 的最后一个维度 d 和 values 的倒数第二个维度 v 进行求和操作。这实际上是一种加权求和，可以看作是注意力机制中的值加权步骤。根据 einsum 的模式 “nshd,nshv->nhdv”，K 的形状是 nshd，即 (1, 4800, 8, 32)。values 的形状是 nshv，即 (1, 4800, 8, 32)。期望的输出形状是 nhdv（n代表批量大小、h代表头数、d代表特征维度、v代表values的最后一个维度），einsum 在 d 和 v 维度上执行了矩阵乘法（而非求和），这意味着 K 的最后两个维度 (d, d) 与 values 的最后两个维度 (v, d) 相乘，得到 (32, 32) 的结果。

然后计算注意力缩放因子Z，因为K本身的形状为（1，4800，8，32），Q的形状为（1，4800，8，32），由于K已经通过.sum(dim=1)在序列长度维度上求和，所以每个头的d维度现在是一个单一值，与Q的d维度进行点积，实质上是对Q的d维度求和，einsum的结果是一个形状为(1, 4800, 8)的张量，表示对每个头和每个序列位置的Q的d维度求和的结果，接下来，对这个结果加上一个小的正数self.eps（防止除零），然后取倒数得到Z。Z的形状也是(1, 4800, 8)。

上述得到Q、KV、Z，后，将其三者进行加权求和，最后得到注意力加权的结果queried_values

下面便是整体的Transformer的q、k、v代码过程，最终输出的结果是一个shape为（1，4800，8，32）的张量，而这个张量具体每一维度的值均是q、k、v通过加权求和计算所得，这便是Transformer的上下文关联，因为q、k、v每个值都有各自的位置编码，序列长度从0~4800，每个序列均有独立的位置编码。

我们回过头来看transformer.py的代码，上述的所有计算过程仅仅是对image0进行处理，transformer.py的前向传播中有两个图像数据输入，因此需要计算两大轮，由于编码器一共有8层，而在进行注意力机制的计算过程中，通过 8个不同的线性变换对 Query、Key 和 Value 进行映射；然后，将不同的 Attention 拼接起来；最后，再进行一次线性变换。每一组注意力用于将输入映射到不同的子表示空间，这使得模型可以在不同子表示空间中关注不同的位置。即每层编码器需要进行8次q、k、v计算。下面的代码是每层注意力的结构，计算完一次q、k、v的值后需要通过全连接层、归一化、前馈神经网络层

通过attention注意力层后，需要经过一个 add&norm 层，它要对输入的向量进行残差连接和 layer normalization层归一化的操作，输出的均是特征维度为256（8*32）的张量。

然后所有的向量都会经过一个 feed forward 层，它本质上就是一个全连接网络，然后得到一排相同维度和数量的向量，再进行残差连接和层标准化操作。

以上便是一层编码器的整体流程，下面代码便是LocalFeatureTransformer的整体网络结构，一共有8层，对于一个输入图像，需要通过8层注意力机制进行计算，其中自注意力和交叉注意力是交错进行，自注意力模块有4层、交叉注意力有4层。

5. 粗粒度特征匹配

接下来回到LoFTR的前向传播中，通过loftr_coarse（Transformer层）后，得到的的feat_c0、feat_c1的各个值均是具有相互的位置关联。然后进行粗粒度特征匹配阶段

CoarseMatching类定义：

self.thr（阈值）和self.border_rm（边界移除）: 这些是一般配置参数，用于调整匹配过程的敏感度或移除图像边缘的干扰。self.train_coarse_percent（设置此目的是为了节省GPU内存）和self.train_pad_num_gt_min（设置此用于避免DDP死锁）: 这些参数用于训练过程中的特定设置。self.match_type: 这是一个配置参数，用于选择不同的可微分匹配方法，在论文和源码中，作者设置的匹配模式是dual_softmax（双softmax）方法，然后设置self.temperature这个参数。

CoarseMatching类前向传播过程：

接下来分析LoFTR粗粒度特征匹配的代码过程，首先使用了map函数和lambda表达式来对两个特征张量feat_c0和feat_c1进行归一化处理，其中map函数接受一个函数和一个可迭代对象作为参数，将函数应用于可迭代对象的每个元素，并返回一个新的迭代器。在这里map函数被用来同时对两个特征张量进行相同的操作。lambda feat: feat / feat.shape[-1]**.5是一个匿名函数（lambda表达式），它接受一个参数feat（在这里代表一个特征张量），并返回feat除以其最后一个维度（通道数）的平方根的结果。这种操作是一种常见的特征归一化方法，可以减少不同特征值范围对后续计算的影响，经过这行代码的处理后，feat_c0和feat_c1将被归一化，使得特征张量中的每个元素的均值接近0，方差接近1，这有助于提高神经网络模型的训练稳定性和收敛速度。

然后使用 torch.einsum 函数计算 feat_c0 和 feat_c1 之间的相似度矩阵 sim_matrix，并除以温度参数 self.temperature 进行归一化。

接着通过应用双softmax策略计算得到的置信度矩阵。这种策略通常用于改善模型对匹配项的识别能力，以获得软相互最近邻匹配的概率。

然后通过get_coarse_match函数得到粗粒度匹配，并将get_coarse_match 函数 所返回的匹配结果和粗粒度关键点坐标值写入data字典中

get_coarse_match函数目的是从置信度矩阵 conf_matrix 中提取粗匹配信息，并根据这些信息更新 data 字典。

首先创建一个名为 axes_lengths 的字典，用于存储不同维度的长度信息。这些信息从 data 字典中获取，分别代表两个特征图的高度和宽度。

接着执行置信度阈值处理，创建一个布尔掩码 mask，其中只有当 conf_matrix 中的元素大于阈值 self.thr 时，掩码才为 True，使用 rearrange 函数重新排列 mask 的形状，将其从二维形式转换为四维形式，以匹配特定的维度顺序

根据 data 字典中是否存在 ‘mask0’ 键，选择不同的处理方式，如果不存在 ‘mask0’，则调用 mask_border 函数处理 mask 的边界。如果存在 ‘mask0’，则调用 mask_border_with_padding 函数，同时考虑已有的掩码信息 data[‘mask0’] 和 data[‘mask1’]，在之前分析的过程中，mask0和mask1均为None。

mask_border函数首先检查边界宽度 b 是否大于0，如果不大于0则不执行任何操作，如果 b 大于0，则对张量 m 的所有维度的前 b 个和后 b 个元素进行遍历，并将其设置为值 v，这种方法通常用于在处理图像或张量时忽略边界区域，例如在图像分割或特征匹配中去除边缘的噪声或不可靠区域。在这里的v在整体的参数配置中设置的为False

最后，再次使用 rearrange 函数将 mask 转换回二维形式，以便于后续的匹配或更新操作

接下来在给定的置信度矩阵 conf_matrix 上执行“相互最近邻”（mutual nearest）匹配的操作，给定 conf_matrix 张量的形状为 (1, 4800, 4800)，这表示它包含一个批次（通常在深度学习中用第一个维度表示），4800 行和 4800 列。conf_matrix.max(dim=2, keepdim=True) 计算 conf_matrix 在第二维度（列）上的最大值。keepdim=True 参数确保结果张量与原张量具有相同的维度数，但被比较的维度上只有一个元素，结果是一个形状为 (1, 4800, 1) 的张量，其中每个元素是原始 conf_matrix 中相应行的最大值。同理conf_matrix.max(dim=1, keepdim=True) 结果是一个形状为 (1, 1, 4800) 的张量，其中每个元素是原始 conf_matrix 中相应列的最大值，然后通过将原始掩码 mask 与上述两个布尔张量进行逐元素乘法（*），只有当一个元素在两个维度上都是最大值时，结果掩码中相应的位置才会是 True。这确保了匹配的互最近邻性质。相当于是找到两幅图像中最有可能匹配上的点。

这种掩码更新方法常用于深度学习中的注意力机制或匹配网络，确保了匹配的对称性和互惠性，即如果特征点 A 将特征点 B 视为最佳匹配，那么 B 也将 A 视为最佳匹配。这有助于过滤掉单向的或不可靠的匹配，提高匹配的质量和准确性。

接着找出两幅图像最有可能匹配上的点，先获取mask矩阵在第二维度为True的点（即第一幅图像最有可能是关键点的点），然后找出这些点的索引值i_ids，根据i_ids和b_ids，得到j_ids，即第二幅图像所能匹配到第一幅图像的点的索引，最后根据i_ids，j_ids，得到置信度矩阵各个点的匹配概率值。

最后4800组图像关键点中一共有1361组关键点可以匹配。

得到关键点的粗略坐标位置后，需要将其映射至原输入图像分辨率的关键点坐标。接着创建一个名为 coarse_matches 的字典，存储用于粗匹配索引和置信度矩阵，随后计算缩放因子scale，它是原始图像高度 hw0_i 与粗特征图高度 hw0_c 的比值。

根据是否存在 scale0 键在 data 字典中，计算第一个特征集的缩放比例 scale0，同理，计算第二个特征集的缩放比例 scale1。‘scale0’ 和 ‘scale1’ 代表两个不同特征集的缩放比例，它们影响匹配点坐标的计算，使得匹配点可以从特征图的分辨率映射回原始图像的分辨率，如果 ‘scale0’ 和 ‘scale1’ 大于1，它们会使匹配点的坐标增大，表示这些点在原始图像中占据更大的位置。如果它们小于1，匹配点的坐标会减小，表示这些点在原始图像中占据较小的位置。

然后将特征图中的匹配点坐标转换为原始图像坐标，torch.stack([i_ids % data[‘hw0_c’][1], i_ids // data[‘hw0_c’][1]], dim=1) 将两个计算得到的索引向量（列索引和行索引）沿着新维度 dim=1 堆叠起来，形成一个矩阵，每一行包含一个匹配点的 (x, y) 坐标（在这里 x 是列索引，y 是行索引）。

最终得到的mkpts0_c、mkpts1_c是一个包含匹配点在原始图像分辨率下的坐标的张量。然后更新coarse_matches

整体的Matching Module模块图：

6. 细粒度特征处理、确定最终关键点坐标

经过粗粒度特征匹配确定匹配的关键点的大致坐标后，接下来对细粒度特征进行特征处理。我们回到 loftr.py 的前向传播代码中：

其中feat_c0、feat_c1、feat_f0、feat_f1 的张量形状为：

而data字典通过粗粒度特征匹配后，其包含了各个关键点索引和置信度矩阵：

接下来，对细粒度特征预处理进行详细分析，fine_preprocess.py（细粒度特征处理和匹配），首先将传入的config配置参数保存为实例变量，从配置字典中获取 fine_concat_coarse_feat 选项，并将其保存为实例变量，这表示是否将粗粒度特征与细粒度特征进行拼接。接着从配置字典中获取 fine_window_size 选项，并将其保存为实例变量，这代表细粒度处理中使用的特征窗口大小，然后从配置字典中获取粗粒度和细粒度模型的维度 d_model，如果配置选项 cat_c_feat 为真，则实例化两个线性层（nn.Linear）。

接下来分析FinePreprocess类的前向传播过程，首先它接受四个特征张量 feat_f0、feat_f1、feat_c0 和 feat_c1，分别代表两组细粒度特征和粗粒度特征，以及一个包含额外数据的字典 data，从类的属性中获取细粒度窗口大小 W，计算步长 stride，这是基于数据字典中原始图像的高度 hw0_f 和粗粒度特征图的高度 hw0_c。

对所有局部区域进行展开（裁剪）操作，以获取固定大小的窗口，然后将特征图中的局部区域提取出来，为细粒度分析或进一步的特征处理做准备，由于feat_f0的张量形状为（1，128，240，320），stride为4，最后得到的特征图大小为 60 * 80，由于是裁剪的每个区域为 5 * 5 * 128，故最后得到的feat_fo_unfold为（1，3200，4800）

如下图，上述代码是将局部区域划分为4800个小区域，这些小区域为55128

然后将其展平，即将55128转换成25*128，这里的 rearrange 调用将局部窗口的顺序从 (n, c, ww, l) 改为 (n, l, ww, c)，其中 n 代表批次大小，c 代表通道数，ww 代表窗口内元素的总数（W**2），l 代表局部窗口的数量，最终得到feat_f0_unfold、feat_f1_unfold，通过 rearrange 函数，局部窗口被重新排列，以便于后续处理。这种排列方式使得每个局部窗口的特征可以更容易地被其他操作（如线性层或自定义处理）访问。

然后接下来对之前所进行粗粒度进行细化，之前粗粒度特征匹配过程只是确定下来的关键点的大致位置，并没有确定下具体的位置，只是将最精确的关键点定位在一个个小范围坐标内。根据粗粒度匹配模块中给出的关键点位置索引来确定关键点在局部窗口的大致位置。我们从下面的debug结果可以看出，一共有1361组关键点可匹配。即这些关键点分别在1361个55128的小区域中。然后微调的过程是对每个5*5的小区域中的25个关键点中选取出最合适的关键点。

上述过程得到了这些关键点分别在1361个55128的小区域中。然后微调的过程是对每个5*5的小区域中的25个关键点中选取出最合适的关键点。

接下来将两幅图像在粗粒度中获取的关键点特征（1361个粗略关键点），即feat_c0、feat_c1通过关键点索引得到1361个关键点，然后两个张量沿着第0维进行拼接，然后，通过调用self.down_proj方法（一个神经网络层或变换函数）对拼接后的张量进行降维或其他变换，得到feat_c_win

其中self.down_proj方法为

【注】d_model_c 为256，d_model_f 为128

得到的feat_c_win的张量形状为：

然后将上述裁剪后的细粒度特征 feat_f0_unfold 和 feat_f1_unfold 沿着第0维进行拼接，拼接后的张量形状为（2722，25，128），2722=1361+1361，调整feat_c_win的形状为（2722，25，128），以匹配 feat_f0_unfold 和 feat_f1_unfold 的形状，最后二者沿着128特征维度进行拼接。粗粒度和细粒度特征融合的结果得到的张量形状为（2722，25，256），最后通过 self.merge_feat = nn.Linear(2*d_model_f, d_model_f, bias=True) 来进行特征降维，将特征维度256降维至128，这便是细粒度特征和粗粒度特征融合的过程

上述便是完成了将粗粒度特征中的关键点信息与细粒度特征中的关键点信息进行了特征融合，接下来将融合后的特征张量feat_cf_win分割成两部分，分别赋值给feat_f0_unfold和feat_f1_unfold。将二者作为LoFTR Module（Transformer）的输入。

通过fine_preprocess.py的特征预处理后，得到feat_f0_unfold和feat_f1_unfold，回到loftr.py的前向传播过程中，进行最后一步，细粒度特征匹配，确定最后的匹配关键点的过程：

由于输入到 Transformer 的数据 feat_f0_unfold 和 feat_f1_unfold 的特征维度均为128维，故细粒度这块的Transformer参数中设置的维度均为128，同时编码器层只有两层，分别是一层自注意力层和一层交叉注意力层。

其余的处理过程均与粗粒度中的Transformer模块思路一致，只不过粗粒度Transformer的输入数据的形状为（1，4800，256），在此细粒度里的输入数据的形状为（1361，25，128），二者的批次数和序列长度、维度均不同，同粗粒度一致，需要先通过线性层生成将128维分成8*16，作为多头注意力层，共有8组特征维度为16的q、k、v进行运算，由于输入的是两组数据，再加上编码器设置为两层，分别是自注意力模块、交叉注意力模块。一共需要4轮计算，最终得到输出结果：

将细粒度特征通过Transformer处理后的数据feat0、feat1作为返回值，赋值给了feat_f0_unfold、feat_f1_unfold，并将其作为最后细粒度特征匹配模块的输入数据。

上述裁剪+特征融合+Transformer处理的架构图：

接下来确定细粒度特征匹配，最终得到精确关键点坐标，首先将Transformer输出的feat_f0_unfold和feat_f1_unfold作为最终输入数据。

接下来计算两个特征张量 feat_f0（feat_f0_unfold） 和 feat_f1（feat_f1_unfold） 之间的相似性，并生成一个热图（heatmap），表示它们在不同位置的匹配概率。

首先从 feat_f0 张量中选择每个批次（batch）的中心特征（一个点）。WW//2 表示沿着特征维度的中间位置，feat_f0 的形状是 [M, WW, C]，其中 M 是批次大小，WW 是窗口大小（可能是特征图的边长），C 是通道数。

然后使用 torch.einsum 计算 feat_f0_picked 和 feat_f1 之间的相似性。einsum 是一个强大的函数，用于根据指定的求和约定执行张量操作。这里，‘mc,mrc->mr’ 表示 feat_f0_picked 的第0维（批次）和第1维（通道）与 feat_f1 的第0维和第2维进行点积，结果的第1维是 feat_f1 的第1维（窗口位置），新张量的形状将是 [M, R]，其中 R 是 feat_f1 的第1维大小。

【注】该操作的主要目的是为了让 feat_f0 张量中选择每个批次（batch）的中心特征（一个点），中心特征的特征维度128，为与 feat_f1 张量中每个批次的 25 个点，每个批次的25个点的维度也为128，进行点积（也就是 $(1,128)\ast (25,128{)}^T$），最后得到的是 每批次feat_f0的1个中心点与 feat_f1 的25个点的相似度，根据相似度大小来判断两个点的匹配度

下面是一个简单的例子，假设 feat_f0_picked 张量形状为（3，10），feat_f1 张量形状为（3，4，10），最后得到的 sim_matrix 张量形状为（3，4）

假设的 feat_f0_picked 张量

假设的 feat_f1 张量

这个小例子里得到的 sim_matrix 张量

得到sim_matrix相似度为：

然后计算 softmax 操作的温度参数（temperature），用于控制 softmax 分布的平滑程度。温度参数是通道数 C 的平方根的倒数，这是一种常见的技术，用于在特征空间中引入一些噪声，防止过拟合。

接着将相似性矩阵 sim_matrix 与温度参数相乘，然后沿第1维（匹配位置）应用 softmax 函数，将相似性转换为概率分布。最后，使用 view 方法将概率分布重新排列成 (M, W, W) 的形状，其中 W 是窗口大小的一半（因为之前选择了中心特征），这样 heatmap 就表示了每个批次中每个位置的匹配概率。

上述过程在原论文中对应的网络结构图：

上述操作的目的是通过特征点积和 softmax 归一化，生成一个热图，用于后续的坐标估计和匹配点的精细调整。接下来计算从热图（heatmap）派生出的坐标及其标准差，并用于后续的细粒度匹配。

然后计算归一化坐标，使用 dsnt.spatial_expectation2d 函数来从热图中获取匹配点的归一化坐标。heatmap[None] 将热图增加一个维度，以便符合函数的输入要求。True 参数表示函数返回的是归一化坐标。结果 coords_normalized 的形状是 [M, 2]，其中 M 是批次中的匹配点数量，2 表示每个点的 x 和 y 坐标。

接着创建归一化网络，创建了一个与热图大小相同的网格，其中 W 是网格的宽度和高度。create_meshgrid 函数（可能是自定义函数）生成了一个二维网格，然后使用 reshape 方法将其变为形状 [1, WW, 2] 的张量，其中 WW 是 W*W，代表网格中点的总数。

随后计算方差，首先计算了网格上每个点的平方，然后与热图相乘，得到加权平方和。接着，从这个加权平方和中减去归一化坐标的平方，得到方差 var。

计算方差，首先对方差进行开方，然后通过 torch.clamp 函数确保数值稳定性（避免因为数值太小而导致的计算问题），最后沿着第二个维度求和，得到每个匹配点的标准差 std。

然后更新数据字典，将归一化坐标和标准差合并到一个新的张量 expec_f 中，形状为 [M, 3]。这个张量被用于细粒度匹配的监督信号，其中最后一位可能表示每个匹配点的置信度或不确定性。

最后计算绝对关键点坐标，调用 get_fine_match 方法，传入归一化坐标和数据字典。这个方法负责将归一化坐标转换为实际的像素坐标，并更新数据字典。

get_fine_match 方法，这个方法的作用是将归一化坐标转换为实际的精细匹配点坐标，并更新数据字典。首先从类的实例变量中提取窗口大小 W，窗口大小的平方 WW，通道数 C，以及缩放比例 scale。然后从数据字典 data 中获取第一个图像的粗匹配点坐标。接着计算缩放因子，随后计算精细匹配点坐标，首先将归一化坐标 coords_normed 乘以窗口大小的一半（W // 2），然后乘以缩放因子 scale1，得到每个匹配点的偏移量。然后，将这个偏移量加到第二个图像的粗匹配点坐标 data[‘mkpts1_c’] 上，得到精细匹配点坐标 mkpts1_f。最后，通过切片 [:len(data[‘mconf’])] 确保坐标数量与置信度数组 mconf 的长度一致。使用 update 方法更新 data 字典，将精细匹配点坐标 mkpts0_f 和 mkpts1_f 添加到字典中，以供后续使用。