wplCZII 第32名解决方案
比赛概述
蛋白质复合物（如携氧的血红蛋白、头发中的角蛋白等）对于细胞功能至关重要，理解它们的相互作用对于人类健康和疾病治疗具有重要意义。冷冻电子断层扫描（Cryo-electron tomography，cryoET）可以生成近原子级别的三维图像（称为断层图），展示蛋白质在其复杂且密集的自然环境中的状态。因此，cryoET 具有解开细胞奥秘的巨大潜力。
目前存在大量的 cryoET 断层图数据尚未被充分利用。越来越多的这类数据以标准格式存储在 cryoET 数据门户网站（cryoetdataportal.czscience.com）中。挖掘这些数据需要自动识别图像中的每个蛋白质分子，即便是人眼可以识别的蛋白质，这一问题仍未得到有效解决。一个可泛化的解决方案将有助于揭示细胞的“暗物质”，推动成千上万的科学发现，并最终促进人类健康。
本次竞赛的挑战在于开发机器学习算法，能够在经过筛选的真实世界 cryoET 数据集中自动标注五类蛋白质复合物。
模型
我们采用了 MONAI 提供的基本 UNet3D 作为主要模型，并使用 Torch 实现了 DLinkNet3D 作为辅助模型。由于训练样本数量有限，模型可能会面临泛化能力不足的问题。为了解决这一挑战，我们设计了一个记忆模块，以增强模型的泛化能力。该模块的具体设计如下：
该模块被添加到 UNet 模型的底层，以增强高维向量的泛化能力。
最终使用的模型如下：
| 模型名称 | 数量 |
|---|---|
| MemUNet | 3 |
| UNet | 4 |
| DLinkNet | 1 |
参数设置：
通道数：(48, 64, 80, 80)
步长：(2, 2, 1)
残差单元数：2
Dropout：0.3
记忆向量大小：(256,256)训练
我们将原始数据的半径缩放至 0.48 或 0.5（感谢 @miyamotodaiya 进行的实验）。在测试了不同的训练尺寸（如 96、128、136、144、164 和 176）后，我们发现 128 和 164 的效果最佳。此外，@yingpengchen 测试了不同的 xyz 尺寸组合，(48, 256, 256) 取得了最优效果。对于输出通道，我们尝试了 6、7 和 8 个通道，最终发现 6 通道的性能最好。
数据增强：
我们使用了以下数据增强方法：
- RanRandCropByLabelClassesd
- RandFlipd
- RandRotated
- RandAffined
- RandGridDistortiond
- RandCoarseDropoutd
- RandScaleIntensityd
- RandShiftIntensityd
我们将旋转和翻转的概率设为 1，以确保数据多样性。
优化器：
我们使用了 @miyamotodaiya 引入的 schedulefree.AdamWScheduleFree 优化器。
损失函数：
- 加权 Tversky 损失
- 距离损失：此损失函数在对真实标签进行距离变换后计算 MSE 损失，使得模型更加关注标签的中心区域。当存在标签重叠时，该损失函数表现尤为优异。
EMA：
我们采用了一种动态调整衰减参数的方法，该方法基于当前模型得分与最佳得分的比较进行调整。这一策略使本地测试结果提升了约 0.001。
推理
在推理过程中，我们使用了滑动窗口策略。对于训练尺寸为 96 的模型，我们在推理时使用了 128（因为 DLinkNet 模型过大，无法处理更大的推理尺寸），而其他模型的推理尺寸为 176 或 180。窗口重叠率设置为 0.15 或 0.5。
推理时间优化：
我们采用了两种推理策略：
- 多模型滑动窗口推理： 该方法结合多个模型进行滑动窗口推理，最终得分 lb 763。我们将模型分布在两块 GPU 上，使用多进程和 TensorRT 加速，并行运行多个模型。共 7 个模型，推理耗时约 4 小时。
- 较少模型 + 滑动窗口 + 大量 TTA： 该策略使用较少的模型，并结合滑动窗口和大量测试时增强（TTA），如翻转和旋转。最终得分 lb 756。所有模型加载至两块 GPU 上，并将数据分为两部分，使用多进程和 TensorRT 加速并行推理两组数据。使用 1 个模型和 7 种 TTA 方法，推理耗时约 5 小时。
最终，我们使用 @miyamotodaiya 提供的数据融合策略，将两种策略的结果合并，最终得分 lb 768。
整体技术流程图
代码
训练代码：https://github.com/wplll/Memory-Enhanced-3D-Segmentation
多模型推理：https://www.kaggle.com/code/peilwang/czii-infer-multi-model
较少模型 + 大量 TTA：https://www.kaggle.com/code/peilwang/czii-infer-fewer-models-and-extensive-tta