通过-Dynamic-SOLO--SOLOv2--和-TensorFlow-更好地理解计算机视觉

通过 Dynamic SOLO (SOLOv2) 和 TensorFlow 更好地理解计算机视觉

原文：towardsdatascience.com/dynamic-solo-solov2-with-tensorflow-for-better-understanding-computer-vision/

github.com/syrax90/dynamic-solov2-tensorflow2 – 文章中描述的项目源代码。

免责声明

⚠️ 首先请注意，这个项目不是生产就绪的代码。

关于项目以及为什么我决定从头开始实现它

该项目针对没有高性能硬件（尤其是 GPU）但想学习计算机视觉或至少在这个领域找到自己的人。我尽量使代码尽可能清晰，因此我使用了 Google 的描述风格来描述所有方法和类，代码内部的注释使逻辑和计算更清晰，并使用单一责任原则和其他面向对象原则使代码更易于阅读。

如文章标题所示，我决定从头开始实现 Dynamic SOLO，以深入理解实现此类模型的所有复杂性，包括整个功能生产的周期，以便更好地理解在计算机视觉任务中可能遇到的问题，并获取使用 TensorFlow 创建计算机视觉模型的有价值经验。展望未来，我要说，我对这个选择没有犯错误，因为它给我带来了很多新的技能和知识。

我会建议所有想要更深入理解其工作原理的人从头开始实现模型。这就是为什么：

当你遇到对某事的误解时，你开始深入探究具体问题。通过探索问题，你找到答案，了解为什么发明了特定的方法，从而扩展你在该领域的知识。
当你理解了某个方法或原则背后的理论时，你开始探索如何使用现有的技术工具来实现它。这样，你就能提高解决特定问题的技术技能。
当你从头开始实现某事时，你会更好地理解可以投入到此类任务中的努力、时间和资源的价值。通过与类似任务进行比较，你可以更准确地估计成本，并对类似工作的价值有更好的认识，包括准备、研究、技术实现甚至文档。

选择 TensorFlow 作为框架仅仅是因为我大多数机器学习任务都使用这个框架（这里没有什么特别的）。

该项目代表使用 TensorFlow2 框架实现 Dynamic SOLO (SOLOv2) 模型。

SOLO: A Simple Framework for Instance Segmentation*，

王新龙，张汝帆，沈春华，孔涛，李磊

arXiv 预印本 (arXiv:2106.15947)*

Dynamic SOLO 图。图片由作者提供。灵感来源于 arXiv:2106.15947

SOLO (通过位置分割对象) 是一个为计算机视觉任务设计的模型，特别是用于实例分割。它是一个完全无锚框的框架，可以预测掩码而不需要任何边界框。论文介绍了该模型的几个变体：Vanilla SOLO，Decoupled SOLO，Dynamic SOLO，Decoupled Dynamic SOLO。实际上，我首先实现了 Vanilla SOLO，因为它是最简单的。但我不会发布代码，因为从实现角度来看，Vanilla 和 Dynamic SOLO 之间没有太大区别。

模型

实际上，根据 SOLO 论文 中描述的原则，该模型可以非常灵活：从 FPN 层数到层的参数数量。我决定从最简单的实现开始。模型的基本思想是将整个图像划分为单元格，其中每个网格单元格只能代表一个实例：确定的类别 + 分割掩码。

SOLO 模型的视觉表示。图片由作者提供。灵感来源于 arXiv:2106.15947

骨干

我选择 ResNet50 作为骨干网络，因为它是一个轻量级的网络，非常适合初学者。我没有使用 ResNet50 的预训练参数，因为我正在实验的不仅仅是原始 COCO 数据集。然而，如果你打算使用原始 COCO 数据集，你可以使用预训练参数，因为它可以节省时间，加快训练过程，并提高性能。

backbone = ResNet50(weights='imagenet', include_top=False, input_shape=input_shape)
backbone.trainable = False

颈部

FPN (特征金字塔网络) 被用作提取多尺度特征的颈部。在 FPN 中，我们使用 ResNet50 对应残差块的所有输出 C2, C3, C4, C5，如 FPN 论文 (由 Tsung-Yi Lin, Piotr Dollár, Ross Girshick, Kaiming He, Bharath Hariharan, Serge Belongie 撰写的特征金字塔网络用于目标检测) 中所述。FPN 的每个级别代表一个特定的尺度，并具有其自己的网格，如上图所示。

注意：如果你使用的是一个小型自定义数据集，其中所有对象的大致尺度相同，你不应该使用所有 FPN 级别。否则，你会训练一些不使用的额外参数，这会导致需要更多的 GPU 资源而徒劳。在这种情况下，你必须调整数据集，使其只为 1 个尺度返回目标，而不是所有 4 个尺度。

头部

FPN 层的输出被用作确定实例类别及其掩码的层的输入。头部包含两个并行分支以实现目标：分类分支和掩码核分支。

注意：根据纯粹的头部架构，我将掩码特征从头部中排除。掩码特征将在下面单独描述。

纯粹的头部架构。图片由作者提供。灵感来源于 arXiv:2106.15947

分类分支（在上述图中标记为“类别”）- 负责预测图像中每个实例（网格单元）的类别。它由一系列按行排列的 Conv2D -> GroupNorm -> ReLU 集组成。我应用了 4 个这样的集。
掩码分支（在上述图中标记为“掩码”）- 这里有一个关键细微差别：与纯粹的 SOLO 模型不同，它不会直接生成掩码。相反，它预测一个掩码核（在论文的第 3.2.3 节“动态 SOLO”中称为“掩码核”），随后通过动态卷积与下面描述的掩码特征应用。这种设计通过减少参数数量并创建一个更高效、轻量级的架构来区分动态 SOLO 和纯粹的 SOLO。掩码分支使用与分类分支相同的结构预测每个实例（网格单元）的掩码核：一系列按行排列的 Conv2D -> GroupNorm -> ReLU 集组成。我在模型中也实现了 4 个这样的集。

注意：对于小型自定义数据集，您甚至可以使用 1 个这样的集来为掩码和分类分支训练，从而避免训练不必要的参数

掩码特征

掩码特征分支与掩码核分支结合以确定最终的预测掩码。该层融合多级 FPN 特征以生成统一的掩码特征图。论文的作者评估了两种实现掩码特征分支的方法：为每个 FPN 级别指定特定的掩码特征或为所有 FPN 级别提供一个统一的掩码特征。像我一样，我选择了后者。掩码特征分支和掩码核分支通过动态卷积操作结合。

数据集

我选择使用 COCO 数据集格式，在原始COCO 数据集和结构相同的小型自定义数据集上训练我的模型。我选择 COCO 格式，因为它已经被广泛研究，这使得编写解析该格式的代码变得容易得多。此外，我选择的LabelMe工具能够将数据集直接转换为 COCO 格式。此外，从一个小型自定义数据集开始可以减少训练时间并简化开发过程。创建数据集的另一个原因是更好地理解数据集创建过程，直接参与其中，并在与 LabelMe 等工具交互中获得新技能。如果你想要更深入地了解 COCO 格式，一个小型的标注文件比大型文件更容易探索和操作。

这里有一些在实施项目过程中遇到的数据集相关子任务（它们在项目中有所展示）：

数据增强。图像数据集的数据增强是通过应用各种图像变换方法来扩展数据集的过程，从而生成与原始样本不同的新样本。掌握增强技术至关重要，尤其是在小型数据集的情况下。我应用了如水平翻转、亮度调整、随机缩放、随机裁剪等方法，以展示如何进行此操作并理解匹配修改后图像的新（增强）图像掩码的重要性。
转换为目标。SOLO 模型期望目标具有特定的数据格式。它接受一个归一化的图像作为输入，没有特别之处。但对于目标，模型期望更复杂的数据：
- 我们必须为每个尺度构建一个网格，通过特定尺度的网格单元格数量来分隔。这意味着如果我们有 4 个 FPN 级别 – P2、P3、P4、P5 – 用于不同的尺度，那么我们将有 4 个网格，每个尺度都有一定数量的单元格。
- 对于每个实例，我们必须通过位置定义它所属的单元格，该单元格是所有网格中的一个。
- 对于每个定义的实例，应用相应的类别和掩码。此外，还有一个问题是将 COCO 格式的掩码转换为只包含掩码像素为 1，其余像素为 0 的掩码。
- 将上述所有内容合并成一个张量列表作为目标。我明白 TensorFlow 更倾向于使用严格的张量集合而不是列表这样的结构，但我决定选择列表，以便在需要改变尺度数量时提供额外的灵活性。
数据集在内存中或生成。数据集分配有两种主要选项：将样本存储在内存中或在运行时生成数据。尽管在内存中分配有很多优点，并且对于很多人来说，将整个训练数据集目录（例如 COCO 数据集）上传到内存（仅 19.3 GB）没有问题——但我故意选择使用 tf.data.Dataset.from_generator 动态生成数据集。原因如下：我认为学习你可能会遇到与大数据交互的问题以及如何解决这些问题是一个很好的技能。因为当处理现实世界问题时，数据集可能不仅包含比 COCO 数据集更多的样本，而且它们的分辨率也可能高得多。使用动态生成的数据集通常在实现上更复杂一些，但更加灵活。当然，如果你愿意，可以用 tf.data.Dataset.from_tensor_slices 来替换它。

训练过程

损失函数

SOLO 没有一个在 TensorFlow 中原生实现的标准化损失函数，所以我自行实现了它。

\[L = L_{cate} + \lambda L_{mask} \]

其中：

(L_{cate}) 是用于语义类别分类的传统 Focal 损失。
(L_{mask}) 是掩码预测的损失。
(\lambda) 系数，在论文中设置为 3。

\[ L_{mask} = \frac{1}{N_{pos}} \sum_k \mathbb{1}_{\{p^*_{i,j} > 0\}} d_{mask}(m_k, m^*_k) \]

其中：

(N_{pos}) 是正样本的数量。
(d_{mask}) 实现为 Dice 损失。
( i = \lfloor k/S \rfloor ), ( j = k \mod S ) — 网格单元的索引，从左到右和从上到下索引。
1 是指示函数，当 (p^*_{i,j} > 0) 时为 1，否则为 0。

\[L_{Dice}=1 – D(p, q) \]

其中 D 是 dice 系数，定义为

\[ D(p, q) = \frac {2 \sum_{x,y} (p_{x,y} \cdot q_{x,y})} {\sum_{x,y} p²_{x,y} + \sum_{x,y} q²_{x,y}} \]

其中 (p_{x,y}), (q_{x,y}) 是预测掩码 p 和真实掩码 q 在 (x,y) 处的像素值。损失函数的所有细节都在 原始 SOLO 论文 的 3.3.2 损失函数 中描述。

从检查点恢复。

如果你使用低性能的 GPU，可能会遇到在单次运行中训练整个模型不切实际的情况。为了避免丢失训练好的权重并继续执行训练过程，本项目提供了一个从检查点恢复的系统。它允许你每 n 个 epoch（n 可配置）保存你的模型，并在以后恢复训练。要启用此功能，将 load_previous_model 设置为 True 并在 config.py 中指定 model_path。

self.load_previous_model = True
self.model_path = './weights/coco_epoch00000001.keras'

评估过程

为了看到你的模型训练得有多有效，以及它在之前未见过的图像上的表现如何，使用评估过程。对于 SOLO 模型，我会将这个过程分解为以下步骤：

加载测试数据集。
准备与模型输入兼容的数据集。
将数据输入到模型中。
抑制相同实例的较低概率的掩码。
显示原始测试图像及其每个实例的最终掩码和预测类别。

在这里我面临的最不规则的任务是实现矩阵非极大值抑制（Matrix NMS），这在原始 SOLO 论文的 3.3.4 节中有描述。NMS 消除了表示相同实例且概率较低的冗余掩码。为了避免多次预测同一实例，我们需要抑制这些重复的掩码。作者提供了 Matrix NMS 的 Python 伪代码，我的一个任务就是解释这段伪代码并使用 TensorFlow 实现它。我的实现：

def matrix_nms(masks, scores, labels, pre_nms_k=500, post_nms_k=100, score_threshold=0.5, sigma=0.5):
    """
    Perform class-wise Matrix NMS on instance masks.

    Parameters:
        masks (tf.Tensor): Tensor of shape (N, H, W) with each mask as a sigmoid probability map (0~1).
        scores (tf.Tensor): Tensor of shape (N,) with confidence scores for each mask.
        labels (tf.Tensor): Tensor of shape (N,) with class labels for each mask (ints).
        pre_nms_k (int): Number of top-scoring masks to keep before applying NMS.
        post_nms_k (int): Number of final masks to keep after NMS.
        score_threshold (float): Score threshold to filter out masks after NMS (default 0.5).
        sigma (float): Sigma value for Gaussian decay.

    Returns:
        tf.Tensor: Tensor of indices of masks kept after suppression.
    """
    # Binarize masks at 0.5 threshold
    seg_masks = tf.cast(masks >= 0.5, dtype=tf.float32)  # shape: (N, H, W)
    mask_sum = tf.reduce_sum(seg_masks, axis=[1, 2])  # shape: (N,)

    # If desired, select top pre_nms_k by score to limit computation
    num_masks = tf.shape(scores)[0]
    if pre_nms_k is not None:
        num_selected = tf.minimum(pre_nms_k, num_masks)
    else:
        num_selected = num_masks
    topk_indices = tf.argsort(scores, direction='DESCENDING')[:num_selected]
    seg_masks = tf.gather(seg_masks, topk_indices)  # select masks by top scores
    labels_sel = tf.gather(labels, topk_indices)
    scores_sel = tf.gather(scores, topk_indices)
    mask_sum_sel = tf.gather(mask_sum, topk_indices)

    # Flatten masks for matrix operations
    N = tf.shape(seg_masks)[0]
    seg_masks_flat = tf.reshape(seg_masks, (N, -1))  # shape: (N, H*W)

    # Compute intersection and IoU matrix (N x N)
    intersection = tf.matmul(seg_masks_flat, seg_masks_flat, transpose_b=True)  # pairwise intersect counts
    # Expand mask areas to full matrices
    mask_sum_matrix = tf.tile(mask_sum_sel[tf.newaxis, :], [N, 1])  # shape: (N, N)
    union = mask_sum_matrix + tf.transpose(mask_sum_matrix) - intersection
    iou = intersection / (union + 1e-6)  # IoU matrix (avoid div-by-zero)
    # Zero out diagonal and lower triangle (keep i<j pairs)
    iou = tf.linalg.band_part(iou, 0, -1) - tf.linalg.band_part(iou, 0, 0)  # upper triangular without diagonal

    # Class-aware IoU: zero out IoU for pairs with different labels
    labels_matrix = tf.tile(labels_sel[tf.newaxis, :], [N, 1])  # each row is labels vector
    same_class = tf.cast(tf.equal(labels_matrix, tf.transpose(labels_matrix)), tf.float32)
    same_class = tf.linalg.band_part(same_class, 0, -1) - tf.linalg.band_part(same_class, 0, 0)
    decay_iou = iou * same_class  # IoU only for same-class pairs (upper tri)

    # Compute max IoU for each mask with any higher-scoring mask
    # (Since i<j is upper tri, for column j, relevant i are those with i < j)
    max_iou_per_col = tf.reduce_max(decay_iou, axis=0)
    comp_matrix = tf.tile(max_iou_per_col[..., tf.newaxis], [1, N])

    decay_matrix = tf.exp(-((decay_iou ** 2 - comp_matrix ** 2) / sigma))

    # Aggregate decay: for each column j, get the minimum decay factor across all i<j
    decay_coeff = tf.reduce_min(decay_matrix, axis=0)  # shape: (N,)
    decay_coeff = tf.where(tf.math.is_inf(decay_coeff), 1.0, decay_coeff)
    # (If no i<j, reduce_min gives +inf; replace inf with 1.0 meaning no suppression)

    # Decay the scores and filter by threshold
    new_scores = scores_sel * decay_coeff
    keep_mask = new_scores >= score_threshold                        # boolean mask of those above threshold
    new_scores = tf.where(keep_mask, new_scores, tf.zeros_like(new_scores))

    # Select top post_nms_k by the decayed scores
    if post_nms_k is not None:
        num_final = tf.minimum(post_nms_k, tf.shape(new_scores)[0])
    else:
        num_final = tf.shape(new_scores)[0]
    final_indices = tf.argsort(new_scores, direction='DESCENDING')[:num_final]
    final_indices = tf.boolean_mask(final_indices, tf.greater(tf.gather(new_scores, final_indices), 0))

    # Map back to original indices
    kept_indices = tf.gather(topk_indices, final_indices)
    return kept_indices

下面是一个例子，展示了模型预测的掩码，这些掩码覆盖了它之前从未见过的图像：

作者提供的带有预测掩码的图像。

从零开始实现的建议

我们将哪些数据映射到哪个函数？ 确保我们给模型提供正确的数据非常重要。数据应该与每一层期望的数据匹配，每一层都会处理输入数据，以便输出适合下一层。因为我们最终是根据这些数据计算损失函数。基于 SOLO 的实现，我意识到一些目标可能不像乍一看那么简单。我在“数据集”章节中描述了这一点。
研究论文。 在你打算基于其构建模型之前，阅读这篇论文是不可避免的。我知道这很明显，但尽管有大量引用其他先前工作和论文，你仍需要理解其原理。当你开始研究一篇论文时，你可能会遇到很多其他论文，你需要在阅读并理解这些论文之后才能进行研究，这可能是一项相当具有挑战性的任务。但通常，即使是最新的论文也是基于一些已经为人所知的原则，这些原则并不是新的。这意味着你可以在互联网上找到很多描述这些原则的材料，这些材料描述得非常清楚。你可以使用 LLM 程序来完成这个任务，这些程序可以总结信息，提供示例，并帮助你理解一些作品和论文。
从小步骤开始。这是一条显而易见的建议，但如果你正在开发一个具有数百万参数的计算机视觉模型，如果你在开发阶段并且不确定模型是否能够正确工作，那么你不需要浪费时间在无用的训练、数据集准备、评估等上。此外，如果你有一个性能较低的 GPU，这个过程甚至会更长。所以，不要从大量数据集、许多参数和一系列层开始。你甚至可以在开发的第一阶段使用小数据集和少量参数让模型过拟合，以确保数据正确匹配到模型的靶标。
调试你的代码。调试代码可以确保你在每个步骤上都有预期的代码行为和数据值。我明白至少开发过一次软件产品的人都知道这一点，他们不需要建议。但我仍然想强调这一点，因为构建模型、编写损失函数、准备与输入和目标交互的数据集，我们经常进行数学运算和张量操作。这需要我们比面对日常的常规编程代码时更多的关注，因为我们知道它的工作原理，即使没有调试。

结论

这是对项目的简要描述，没有技术细节，以提供一个总体印象并避免阅读疲劳。显然，一个专门针对计算机视觉模型的项目的描述不可能在一篇文章中完成。如果我看到读者对项目感兴趣，我可能会写一篇更详细的分析，包括技术细节。

感谢阅读！

posted @ 2026-03-28 10:06 布客飞龙III 阅读(12) 评论(0) 收藏举报

刷新页面返回顶部

龙哥盟