MICCAI 2021 | BiX-NAS：为医学图像分割搜索高效的双向架构

点击下方卡片，关注“CVer”公众号

AI/CV重磅干货，第一时间送达

导读

如何轻量化模型？本文提出了一种高效的NAS方法，搜索出了超轻量的分割模型，在保证准确率的情况下大大节约了计算成本。

BiX-NAS: Searching Efficient Bi-directional Architecture for Medical Image Segmentation

作者单位：悉尼大学，新南威尔士大学，匹兹堡大学，京东金融

项目网页（含代码以及论文）：

https://bionets.github.io/

论文：https://arxiv.org/abs/2106.14033

代码：https://github.com/tiangexiang/BiX-NAS

最近，可循环的U-Net网络变体已被引入到各种医学图像分割任务中。现有的研究侧重于通过重用构建块来实现网络循环，在不显著增加参数的同时提升网络的性能。但这种可循环的模块仍然不可避免地增加了计算成本。

本文研究了双向跳跃连接 (bi-directional skip connection) 网络的多尺度（multi-scale) 变形，然后通过一种新颖的两阶段式神经结构搜索 (NAS) 算法BiX-NAS，自动搜索出一种超轻量且有效的可循环结构。本文提出的方法通过过滤掉无效的多尺度跳跃连接来降低网络的计算成本。在三个不同的医学图像数据集上的实验表明，BiX-NAS搜索出的结构（BiX-Net）以显着降低的计算成本实现了最先进的性能。

方法简介

1. BiO-Net++：轻量的BiO-Net多尺度升级版

如何轻量化模型是一个非常重要的问题。在一个前序工作中，BiO-Net会触发多个编码（encoding) 和解码 (decoding) 阶段，将同一尺度的特征用双向跳跃连接融合起来（图1左）。尽管这种循环设计可以大大缩小网络规模，但根据其预设的迭代次数，计算成本仍然不可避免地增加。

本文首先将多尺度的跳跃连接引入BiO-Net，提出了BiO-Net++（图1中）。注意，特征融合策略由concatnation优化为element-wise平均，以降低网络总体复杂度。

虽然上述设计促进了多尺度特征融合并缩小了网络规模，但从实验中发现，这种密集的连接会增加过量的计算量却仅仅在整体性能方面带来微小改善（表1）。本文目标寻找一种更稀疏连接的BiO-Net++子结构，让模型不仅可以从多尺度融合中受益，还可以最大程度地减轻计算负担。

2. BiX-NAS: 超高效搜索BiO-Net++子结构

为此，本文提出了一种两阶段式搜索算法BiX-NAS，以找到稀疏连接的子结构。在阶段一中，其采用了一种可训练的selection matrix来缩小整体的搜索空间，并在阶段二中引入Evolutionarty NAS来逐步进化出最佳的子结构。

2.1. Phase1：通过selection matrix缩小搜索空间

为了降低整体搜索难度，本文采用可微分NAS算法来快速筛掉那些易于发现的无效跳跃连接。

假设待搜索的模块有N个输入跳跃连接，本文预设在稀疏连接的结构中，只有 k ∈ [1, N - 2] 个候选的跳跃连接可以被保留。因此，对于有L层和T次迭代的BiO-Net++ 来说, 搜索空间大概是：

本文通过构建一个可学习的selection matrix M来对N个传入连接和k个候选连接之间的映射Φ(·)进行建模，并将Φ(·)公式化为以下完全可微的方程：  

2.2. Phase2: 渐进式进化搜索确定最佳子结构

为了进一步减少阶段一中搜到结构的潜在冗余，本文引入了一种新颖的进化搜索算法以找到更好的网络。

具体来说，此算法将同时搜索一对extraction stage之间所有的跳跃连接 （见上图黄色标注），然后在当前搜索结束后逐步移动到下一对。注意，由于网络依赖从前往后的拓扑顺序， 因此在搜索时需要从最后一对extraction stage开始，并逐步移动到第一个。

在传统的进化NAS方法中，最直接的策略是从SuperNet采样不同的子结构（由某些跳跃连接组成），然后分别进行独立的训练。这种传统策略有两个主要缺陷：第一，分开训练子结构可能会导致不公平的结果；其次，根据实验和经验，搜索过程非常缓慢且低效。

2.3. 分析搜索的公平性

为了克服上述第一个缺陷，本文提出了skip fairness的概念，并声称所有连接搜索算法都需要满足这一原则。请注意，每个采样的子结构都是从不断更新的SuperNet中随机抽取的，它构成了每次迭代的总体P。

上述概念要求，在不同的采样子结构中（e.g. ），任何对应的跳跃连接 （e.g. ）需要携带相同特征。否则，不一致的特征会影响搜索决策，从而导致不公平结果。由于单独训练采样的子结构产生的传入特征不一致，传统的策略违反了这一原则，因此不直接适用于搜索跳跃链接。

本文提出的阶段二NAS算法通过共享部分的传入特征来满足skip fairness。具体来说，假设在第t 和第t+1个extraction stage之间搜索时：从第1个到第t-1个extraction stage的结构是固定的，且可以共享这些stages的生成特征。这些stages组成的子结构被称为head network （图2a）。相反，第t到最后一个extraction stage的网络结构随着不同的采样而变化，这种不固定的stages组成的子结构集被称为tail networks（图2a)。这些tail networks共享相同的超网络权重，但具有不同的连接方式。如图2b，head network生成的特征被传入到所有采样的tail networks中。然后对所有tail networks的loss进行平均，并且仅计算gradients一次。因此，阶段二的实际搜索过程非常的高效，克服了第二个缺陷。

在每对提取阶段之间的搜索完成后，如图2c，本文遵循多目标选择标准，该标准基于验证准确性 (IoU) 和计算复杂性 (MACs) 保留在帕累托前沿的子结构。算法细节和搜索到的BiX-Net详见论文。

实验结果

BiX-Net在nuclei segmentation和organ segmentation上都达到了SOTA的准确率，但是仅需0.28M参数量和28G MACs。U-Net的计算量超过BiX-Net将近23倍，BiO-Net的计算量则是超过将近40倍。  

论文PDF和代码下载

后台回复：BiX，即可下载上述论文和代码

后台回复：ICCV2021，即可下载ICCV 2021论文和代码开源的论文合集

后台回复：CVPR2021，即可下载CVPR 2021论文和代码开源的论文合集

CVer-医学图像交流群成立

扫码添加CVer助手，可申请加入CVer-医学图像 微信交流群，方向已涵盖：目标检测、图像分割、目标跟踪、人脸检测&识别、OCR、姿态估计、超分辨率、SLAM、医疗影像、Re-ID、GAN、NAS、深度估计、自动驾驶、强化学习、车道线检测、模型剪枝&压缩、去噪、去雾、去雨、风格迁移、遥感图像、行为识别、视频理解、图像融合、图像检索、论文投稿&交流、PyTorch和TensorFlow等群。

一定要备注：研究方向+地点+学校/公司+昵称（如医学图像+上海+上交+卡卡），根据格式备注，可更快被通过且邀请进群

▲长按加小助手微信，进交流群

▲点击上方卡片，关注CVer公众号

整理不易，请点赞和在看