第二名 - nnU-Net + blob 回归

非常感谢 @andrewjdarley、BYU 和 Kaggle 举办此次竞赛！也要大声感谢 @brendanartley 慷慨地与所有人分享他的外部数据！

概述 (TLDR)

以下是我们解决方案的简要 rundown —— 它 straightforward 且易于实现：

• 我们将电机定位任务 formulate 为 blob 回归，使用 TopK (20%) BCE 损失进行优化。
• 我们基于 nnU-Net 构建，这是领先的 3D 医学图像分割框架。通过一些简单的技巧，它可以适应 blob 回归。
• 我们的模型是一个带有残差编码器的 3D U-Net，从头开始训练。
• 我们使用了竞赛数据、Bartley 的数据以及 555 个额外的公开断层图。电机标注经过了手动修正。
• 推理使用单个模型和轻量级的测试时增强（镜像）。

我们的模型在公共/私人排行榜上取得了 0.86734 / 0.87656 的分数，与第一名并列。Kaggle 按提交时间解决并列情况。这对我们来说很不幸——但 Bartley 确实当之无愧。恭喜！

我们是谁？

我们是一群同事（科学家和博士生），隶属于德国癌症研究中心的 医学图像计算部 和 智能医疗系统部，以及 Helmholtz Imaging。我们的专业知识在于 3D 图像分析——特别是在解决 3D 分割问题和开发将算法引入临床的基础设施方面。对于我们 5 人中的 4 人来说，这是第一次认真参加 Kaggle 竞赛，尽管我们在医学图像分割挑战方面有过往绩。

使用的数据

我们使用了竞赛数据 (n=648)，Bartley 的外部数据 (n=1287) 以及另外 555 个公开可用的图像（总共 n=2490）。

Bartley 的数据。 我们没有直接使用他提供的数据，而是使用他提供的 CziiCollector 的修改版本重新下载所有图像。这样做有两个原因：a) 他的调整大小策略与我们的方法不同（见预处理），要求我们从原始断层图开始；b) 他只使用了下载数据的一部分（104 个数据集中的 62 个和 3395 个断层图中的 1287 个）。

修正。 我们使用早期模型的 5 折交叉验证预测来生成所有官方断层图和 Bartley 数据 (n=1935) 的预测。我们配置了一个非常低的检测阈值以减少 FN，代价是增加 FP。我们将 GT 和预测编码为实例分割图，其中球体代表电机，然后使用 napari 数据检查工具 手动检查所有断层图的标注错误。修正了 213 个断层图，最常见的错误是：

• 在许多电机实例的图像中 missing motors
• 在看不见鞭毛的图像中错误标注的电机
• 靠近图像边缘 missing 的电机。

我们强调，在本次竞赛之前，我们团队成员都没有深入了解 cryoET 或细菌形态的表现。在整个挑战过程中，我们使用提供的训练数据、ChatGPT 和 Google 图片搜索训练了我们的“生物神经网络”。因此，虽然我们做出了 genuine 的努力，但我们的手动修正可能并不完全准确。

额外数据。 如上所述，Bartley 的数据集仅涵盖了他的 CziiCollector 下载的断层图的一部分。注意下载的数据结构化为数据集 (n=104)。我们使用我们的一个模型（公共分数约 0.85/0.86）预测所有尚未包含在 Bartley 集合中的断层图。我们使用以下策略采样额外数据：

• 在所有预测有电机的断层图中（约 500 个），随机采样 250 个
• 对于每个数据集，随机采样 4 个额外的没有电机的断层图（如果可用断层图较少则更少）。
• 最后，通过确保每个数据集至少包含 4 个含电机的断层图来增加电机外观多样性（如果数据集有这么多，大多数没有！）

我们的额外数据包含 555 个额外的断层图。这 555 个案例随后使用上述相同策略进行了手动修正。这将训练案例的数量从 1935 个增加到了 2490 个。

在合并外部数据时，我们模拟竞赛对外部数据的强度处理，将其裁剪到 0.1 和 99.9 百分位数并转换为 uint8，类似于 Bartley 的做法。这只是为了使数据与官方数据（和测试集）保持一致。

预处理

由于测试集不提供体素间距（这本来是更可取的），我们决定调整所有断层图的大小，使最长边为 512 像素。
nnU-Net 自动执行 z-score 归一化。每个图像都转换为 float32，并通过减去其均值并除以其标准差进行单独归一化。由于原始数据是 uint8，这一步可能有些浪费，但我们没有时间调查其他策略，坦白说，也没看到这里有什么改进的潜力。

网络架构

我们使用 nnU-Net 的 ResEnc，它本质上是一个带有残差编码器和轻量级卷积解码器的 UNet。我们也试验了 nnU-Net 的标准 UNet，老实说，没有看到显著的性能差异。任一模型表现都非常好。

训练过程

我们从头开始训练所有模型。最初，我们进行了 5 折交叉验证以剔除糟糕的设计选择。Blob 回归允许每个图像有多个电机预测，我们开发了一个内部评估方案，兼容任意数量的电机，从而允许我们使用所有训练数据进行内部验证。后来我们依赖排行榜提供反馈，因为我们观察到的内部性能与公共分数之间的差距使我们在最终模型优化时不信任内部结果。

Blob 回归

Blob 回归是 landmark 检测的标准方法，此前已在竞赛中成功应用。一般前提是你可以通过预测 blob 的局部最大值恢复对象定位。

我们如何使用 nnU-Net 进行回归

nnU-Net 是为语义分割构建的。这也包括其预期的数据结构。为了使其兼容电机回归，我们将 ground truth 存储为实例分割图，其中每个电机编码为一个球体（r=6 像素），具有唯一的整数标签。这些球体被 nnU-Net 视为分割，并像 nnU-Net 通常那样通过数据加载和增强管道，从而正确应用旋转、镜像等。在数据加载管道的末尾，我们注入一个自定义变换，将每个电机实例转换为 blob。通过像素级 max 操作注入 blobs，以正确允许近距离的电机。

Blob 生成

我们使用'EDT blobs'，基本上是经过欧几里得距离变换并重新缩放为 [0, 1] 值范围的 3D 球体。EDT 球体比高斯球体有更锐利的‘中心’。我们也试验了高斯球体，得到了非常相似的结果。需要更多实验才能给出哪个更好的 definitive 答案。

我们的球体半径为 25 像素。这 eased 学习，因为目标张量中的不平衡（由 0 主导）减轻了。我们在 r=15 时也取得了非常好的结果。同样，我们没有足够的时间进行 intensive 实验。

以下是生成的 ETD blobs 的两个示例：

损失函数

正如其他人报道的那样，MSE 和 soft Dice 损失在此任务上表现不佳。我们仅使用 20% 最差的体素（在整个 batch 上计算的具有最高损失值的体素）计算 binary cross-entropy 损失。这在早期实验中相对于常规 BCE 给出了一个小幅提升，因为它 counteracts 目标中的不平衡。Focal 损失没有优于 TopK20 BCE。

超参数

我们的最终模型训练 batch size 为 16，patch size 为 (128, 256, 256)。初始学习率为 0.01，并使用 polyLR 计划（与默认 nnU-Net 相同）在训练过程中衰减。我们使用 SGD 训练 3500 个 epochs，其中每个 epoch 定义为 250 次迭代。因此，我们的模型训练了 3500*250=875,000 步，并看到了 3500*250*16 = 14,000,000 个 patches。

Patch 采样策略

我们利用 nnU-Net 的默认策略，这与此处使用的实例分割内部表示 synergizes 良好。采样策略如下：

• 对于 batch 中 67% 的样本：
  • 选择一个随机训练案例
  • 从该案例中选择一个随机 patch
• 对于剩余的 33%（每个 batch 至少 1 个样本！）
  • 选择一个随机训练案例
  • 如果此案例有任何电机：
    • 选择一个随机电机实例
    • 选择一个包含此电机的 patch
  • 如果没有电机，回退到随机采样

数据增强

我们在训练期间应用重度数据增强，包括：

• 随机强度裁剪 (10%)
• 旋转和缩放 (各 30%)
• Rot90 (50%)
• 转置兼容轴 (形状相同的轴 = 两个 256 轴，见 patch size, 50%)
• OneOf(MedianFilter (10%), GaussianBlur (10%))
• 高斯噪声 (15%)
• 加法亮度 (5%)
• OneOf(Contrast (15%), Multiplicative Brightness (15%))
• 模拟低分辨率 (7.5%)
• Gamma (20%)
• Gamma (反转图像) (20%)
• 镜像 (每个轴 50%，可以镜像多个轴)
• 加法亮度梯度 (10%)
• 局部 Gamma 梯度 (10%)
• 锐化 (10%)
• 图像反转 (10%)

我们参考 定义变换的代码 获取更多细节，这里太多（且不够 significant）无法全部列出。

数据加载基础设施

其他人报告了 on the fly 加载和增强断层图的问题。得益于 batchgeneratorsv2 快速增强实现和 nnU-Net 高效的数据基础设施，我们没有观察到此类问题。我们使用 blosc2 作为数据格式，允许从压缩文件中部分读取，使我们能够仅读取当前 patch 所需的断层图部分。此外，通过使用 memmap 读取，我们可以有效地将一些读取缓存到 RAM 中（通过 OS 自动完成），从而在从网络驱动器读取数据时减少网络带宽。数据加载和增强在 CPU 上完成。

计算资源需求

我们的最终模型是在 8xA100 40GB 上使用 PyTorch 的 DDP 训练的。训练耗时略少于 7 天。注意我们只在最后阶段 scaled 了计算资源。我们较少计算资源的最佳模型得分为 0.86392（私人 lb），在单个 A100 上训练了约 18 小时。注意比较并不理想，因为我们在较小模型上花费的阈值优化时间要少得多。

阈值调整

在运行内部交叉验证时，我们计算所有电机检测，然后可以有效地 sweep 最佳阈值，允许在各自的 sweet spot 比较模型并确定阈值稳定性。在切换到排行榜进行模型优化后，我们每个模型花费 5-10 次提交来确定最佳阈值，没有遵循系统策略。我们没有做百分位数阈值处理，事后看来，也许应该这样做，以更有效地利用提交次数。
对于我们的最终模型，0.15 对于公共 (0.86734) 和私人 (0.87656) 排行榜都是理想的。‘低’阈值与使用 EDT 而不是高斯 blobs 有关。

推理

我们主要使用 nnU-Net 的推理基础设施。断层图被 dissect 成一系列重叠的 patches（50% 重叠），预测通过加权当前预测 patch 的中心像素高于边界（高斯重要性加权）来 stitch 在一起。测试时增强通过沿所有轴镜像应用。预测的 logits 通过应用 sigmoid 被 clamp 到 [0, 1]，随后进行电机检测。
我们使用 2xT4 实例进行预测，并将工作负载 evenly 分配给两个 GPU。我们始终使用单个模型，没有 ensembling。推理耗时 7-8 小时。

后处理

交叉验证

预测的 blobs 通过执行 non-maximum suppression 转换为电机预测。在交叉验证期间，我们需要允许每个断层图有多个电机检测。我们用 3D 高斯模糊预测的 blobs（可选）。然后我们检测电机为 motors = torch.argwhere((prediction == max_pool(prediction, kernel_size=min_motor_distance)) & (prediction > threshold))。

排行榜

排行榜只有 0 或 1 个电机的断层图，允许我们简化推理逻辑。我们只需找到预测中的最大强度，并检查它是否高于电机检测阈值。

结果

我们的模型（单个 checkpoint，没有 ensembling）取得了 0.86734 的公共分数和 0.87656 的私人分数。这与 Bartley 的解决方案并列。不幸的是，Kaggle 按提交日期解决并列，从而授予 Bartley（ admittedly 非常当之无愧的）第一名。在截止日期前 9 天进行最后一次提交需要相当大的勇气。为此点赞！

我们非常想为最终提交所做的不同设计选择提供适当的 ablations，但觉得这只会误导，因为我们没有为所有模型提供相等的阈值调整预算，并且没有 checkpoints 用于有趣测试场景的相同模型的proper 1:1 比较。也就是说（请持保留态度），以下是一些 anchors（私人分数，报告符合描述的最佳提交）：

数据

低计算资源 (1xA100 40GB, 18h 训练) 比较

• 未修正官方数据：N/A (抱歉)
• 未修正官方 + bartley 数据：0.83181
• 修正官方 + bartley 数据：0.86253
• 修正官方 + bartley 数据和 555 个额外案例：0.86392 (可能低于应有水平，由于阈值优化不足！)

=> 修正 GT 似乎产生了很大影响。额外数据的影响 unclear
高计算资源比较在这里没有意义，因为样本不足且结果杂乱无章。

高斯 vs EDT blobs

低计算资源 (1xA100 40GB, 18h 训练) 比较，使用修正官方 + bartley 数据

• 最佳 EDT: 0.84888 (r=25)
• 最佳高斯：0.84513 (r=15)

对于其他所有内容，我们的数据点不足，阈值调整预算太不平衡，或者实验配置差异太大。

什么方法不起作用？

虽然我们确信 blob 回归是理想的任务 formulate，但我们想通过尝试其他任务 formulate 来 doubly sure：

• 带有后处理的 3D 分割（也是 nnU-Net）
• 基于 YOLO 的 2D 检测
• 基于 nnDetection 的 3D 检测
• 使用 nnLandmark 的 landmark 检测（也做 blob 回归但使用 MSE 损失）

在初始实验中，这些都没有接近基于 nnU-Net 的 blob 回归性能，并很快被 discontinued。注意，每个解决方案都可能进一步优化以达到 competitive 性能——我们只是没有投入更多时间，只是 out of the box 尝试了它们。

其他损失 formulate 如 soft Dice, focal loss, MSE 没有帮助。标准 BCE 与我们在此使用的 TopK 变体性能相似。

我们试验了 FP oversampling，通过增加采样我们先前模型迭代生成 FP 电机预测的 patches 的可能性。这导致了大致等效的性能，并由于额外的复杂性而被丢弃。

我们还应该做什么？

我们加入得晚，早期没有投入足够的时间，所以最后处于时间压力下，并受到提交限制的极大约束。我们绝对应该更早开始，更系统地利用提交次数。

其他人报告 quantile 阈值处理是克服 lb 上阈值优化需求的好解决方案。我们应该这样做。

我们没有投入足够的时间进行 ensembling，导致我们的最终模型是单个 checkpoint。使用 ensembling 可能有一些性能改进。有效地做到这一点需要我们训练更小/更快的模型，并在推理中仔细平衡 ensembling 与 TTA 和 patch 重叠，所以这不是我们一夜之间就能做到的。

我们希望什么？

直到今天，我们仍然不知道是什么导致了内部 CV 和排行榜之间的性能差异。我们怀疑可能存在 distribution shift，例如，电机总数不同，细菌物种不同，或扫描仪不同。不得不如此依赖排行榜感觉相当 frustrating。如果有一个训练数据集允许有意义的内部验证就好了，这样我们可以测试更多想法，并且不受每天 5 次提交的限制。所以，本质上是一个更能代表预期目标分布的训练数据集。

我们发现对于通常以 float32（或偶尔 uint16）运行的模态，使用 uint8-quantized 强度有些 limiting。也不清楚组织者应用了哪些额外的预处理步骤（例如，强度裁剪或归一化），这在集成外部数据时引入了一定程度的 guesswork。虽然我们理解这一选择可能是为了更包容来自计算机视觉社区的参与者，但感觉就像拉着手刹开车。提供 full-precision 数据以及转换为 jpg/png 的脚本将提供两全其美的方案。

调整到共同的体素间距是 3D 图像（如断层图）的标准程序，在这里做 would have been good。我们想知道为什么测试集中没有提供体素间距信息。

训练和测试数据集中似乎存在 已知错误，组织者未予修正。虽然我们理解这会 upset 一些参与者，但我们认为最好更新标注，特别是在私人测试数据集上，以确保我们准确测量算法性能。

致谢

我们感谢 BYU，特别是 Andrew Darley 的组织，以及 Kaggle hosting 此次竞赛。我们还要再次感谢 Bartley，慷慨地在一个他冒着被人使用数据超越其解决方案的风险的环境中分享他的数据——这是一个勇敢的举动。此外，我们要向德国癌症研究中心 (DKFZ) 的医学图像计算部和智能医疗系统部以及 Helmholtz Imaging 致敬，因为它们太棒了。我们还要感谢 Lars Krämer 出色的 napari 数据检查工具，这使得手动检查电机标注变得轻而易举。最后，非常感谢团队——一起参加这次竞赛真是一次 amazing 的经历！