HoVer-UNet申请出战 | 精度不变，速度提升3倍，UNet家族还真是生生不息！！！

435 0 0

点击下方卡片，关注「集智书童」公众号

HoVer-UNet申请出战 | 精度不变，速度提升3倍，UNet家族还真是生生不息！！！

作者提出了”HoVer-UNet”方法，用于提取多分支HoVerNet框架的知识，用于核实例分割和分类在组织病理学中。作者提出了一个紧凑、简洁的单UNet网络，其Backbone为Mix Vision Transformer，并配备了自定义的损失函数，以最优地编码HoVerNet提取的知识，在不降低计算要求的同时，保证性能。

作者展示了在公共Pan-Nuke和Conseq数据集上实现了与HoVerNet相当的结果，推理时间降低了三倍。

代码：https://github.com/DIAGNijmegen/HoVer-UNet

1 Introduction

核子全视野分割，即同时检测、分割和分类核实例，是数字病理学中多个任务自动化的核心，尤其是在常规伊红-复染色（H&E）染色组织切片分析中。近年来，这一任务越来越多地使用深度学习技术来解决，实现了能够驱动计算机生物标志物设计的准确分割结果。然而，已有几项研究针对核子实例分割和分类问题进行了处理，一个关键的关注点是准确结果所需的推理时间。

之后，HoVerNet使用NP和HV进行后处理以获得核子实例，并将其与TP输出相结合以对核子进行分类（图1）。该网络在多个公共数据集上的有效性已经得到证明，但它的尺寸和复杂性导致了高的推理时间。

知识蒸馏。由Hinton等人引入，知识蒸馏通过使用教师网络的预测来指导并改进学生网络的学习过程，从而基于更复杂的教师网络训练一个较小的网络（学生）。为此，它使用一个损失函数，将两个超参数合并：

温度（）用于平滑预测，使学生可以学习，同时减少教师偏见的影响，
alpha（），将学生损失（计算学生预测和真实值之间的差异）与蒸馏损失相结合。

后者被计算为学生和教师标准化softmax logits之间的Kullback-Leibler差异，乘以。在本论文中，作者提出使用KD从原始HoVerNet中衍生出一个轻量级的核子全视野分割模型，目的是在保持HoVerNet性能的同时，加速推理时间。

2 Method

在这个部分，作者详细介绍了HoVer-UNet，重点介绍在训练过程中使用的蒸馏策略，并引入了所提出的损失函数。

蒸馏框架

在作者KD框架中，作者采用一种离线技术，使用HoVerNet作为预训练的教师网络。鉴于HoVerNet通过三个分支对核子实例进行分割和分类，作者的蒸馏策略是基于将HoVerNet的所有输出分支合并为一个分支网络（参见图1.1）。

请注意，作者的目标是训练一个学生模型，可以替换HoVerNet的Backbone网络，而不是其后处理步骤，后者作者保持不变。作者使用单分支UNet作为学生模型，并将HoVerNet的所有分支输出合并为一个具有与HoVerNet分支总数相等输出通道的单分支。特别是，作者使用Mix Vision Transformer（MixViT）作为UNet的Backbone网络，这是根据作者的实验得出的最佳组合。

损失函数

作者提出一个自定义的损失函数，将HoVerNet蒸馏为一个单分支UNet。根据KD理论的建议，作者的损失是一个线性组合，包括学生和真实值之间的学生损失，以及由参数调节的学生和教师之间的蒸馏损失，定义为：

对于单个损失项，作者受到HoVerNet损失的启发，并针对KD进行了具体的修改。具体而言，学生和蒸馏损失，分别表示为和，是从三个分支（NP、HV和TP）上的损失的线性组合导出的。对于每个分支的总损失，定义为这些分支的单个损失之和，其中表示学生或蒸馏情况。

在以下所有公式中，表示学生预测，表示蒸馏损失预测，而表示学生损失预测，即真实值。

对于HV分支，损失被定义为，其中包含均方误差（MSE）和均方梯度误差（MSGE）。
对于NP和TP分支，学生损失被定义为，其中结合了加权交叉熵（CE）损失和Dice损失（DICE）。
与学生损失不同，这些分支的蒸馏损失表示为，利用加权交叉熵（CE）损失和Kullback-Leibler Divergence（KLD）。

3 Experimental Results

在本工作中，作者使用了两个数据集，分别是用于训练HoVer-UNet的Pan-Nuke和用于在外部数据上验证结果的CoNSeP。PanNuke包括6078个伊红-复染色（H&E）瓷砖，每个瓷砖的大小为，覆盖了19种不同的组织类型。该数据集是针对核子实例分类和分割设计的，并包含了五种核子类型的注释：恶性、炎症、连接/软组织、死亡和健康上皮。

CoNSeP包括41张伊红-复染色（H&E）图像瓷砖，每个瓷砖的大小为像素，在目标放大倍数下，带有注释的七种核子类型：恶性、正常、内皮、杂项、成纤维细胞、肌肉和炎症。

为了评估作者框架的有效性，作者使用了Graham等人[4]和Gamper等人[10]研究中推荐使用的指标。具体而言，作者采用了最初在[1]中提出的Panoptic Quality指标，并结合Graham等人[4]引入的F分数。为了训练作者的网络，作者使用了Adam优化器，初始学习率为，而beta值为。

作者还使用了具有耐心为5，缩放因子为，偏差为，最小学习率为的减少平台学习调度器，以及具有耐心为10的早期停止。此外，作者使用了编码器和解码器的深度为5，并将编码器初始化为ImageNet预训练权重。

超参数定义。作者的框架包括三个超参数：温度系数()，参数和Backbone网络。在作者的实验中，作者探索了这些参数的各种值。对于，作者检查了1、3和5的值。对于，作者考虑了0和0.5的值。

Backbone网络选择。关于Backbone网络，作者尝试了各种架构，包括MixViT-B0、MixViT-B1、MixViT-B2和MixViT-B3。作者选择了这些模型，因为它们比HoVerNet快得多，且减少了Multiply-Accumulate Operations（MACs）。

具体而言，最复杂的模型拥有10.63 GMac，推理时间为0.065秒，而HoVerNet具有149.73 GMac，推理时间为0.835秒。作者将推理时间计算为每个批次的16个256×256大小的块的平均值。

在PanNuke数据集上，作者评估了包括二进制全视野质量()、多类别全视野质量()、特定类别的全视野质量()、检测F分数()和单个类别的F分数()在内的多个性能指标。

类别分别表示恶性（N）、炎症（I）、连接/软组织（C）、死亡细胞（D）和上皮（E）。为了确保稳健性，作者采用了一种三倍交叉验证（3-CV）方法进行训练，这是Pan-Nuke作者[10]的建议。所有报告的结果都代表测试结果的平均值。

考虑到作者实验中收集的大量数据，作者将最佳性能与DIST、Mask-RCNN、Micro-Net和HoVerNet进行了比较。具体而言，使用MixViT-B2 Backbone网络的UNet模型，将和作为超参数，实现了最佳性能。

Table 1呈现了PQ评估结果。与HoVerNet相比，作者的解决方案获得了较低的分数，但与HoVerNet的性能相符，并超过了表格中列出的其他网络。

然而，当考虑的结果，如表2所示，作者的解决方案与HoVerNet具有可比性能，并优于其他网络。此外，作者的提出的解决方案在处理速度方面相对于HoVerNet具有显著优势。

这主要归因于作者设计的架构旨在减少推理时间。此外，作者保留了与先前概述相同的后处理方法，进一步巩固了作者的方法作为此任务的有效选择。在CoNSeP数据集上，作者进行了HoVerNet和HoVer-UNet之间的比较分析，两者都在PanNuke上进行预训练，并利用CoNSeP数据集对跨域数据进行外部评估。

鉴于核子分类在PanNuke和CoNSeP之间以及训练好的HoVer-UNet和CoNSeP目标之间只有部分对应关系，因此作者为更详细的比较将标签映射到多个子类。新定义的类包括恶性、炎症、上皮和杂项。

Table 3显示了作者的解决方案在PQ方面优于HoVerNet，虽然在F分数检测方面稍逊一筹。关于分类指标，作者的解决方案在恶性核子和上皮核子方面优于HoVerNet；在杂项和炎症核子方面则基本相等，但在炎症核子方面较差。

最后，推理时间约为HoVerNet的1/3。图2展示了与CoNSeP参考标准相比，HoVerNet和HoVer-UNet的结果的视觉示例。总体而言，结果之间的相似性支持了作者方法的实用性。