我的 CZII CryoET 第 43 名解决方案

这不是获胜的解决方案，但我在这里分享我的努力成果，以防任何人感兴趣。

我首先强调，几乎所有的模型训练和大部分编码都是使用 Kaggle 计算完成的，因为我的家用计算机不足以支持进一步的探索。

模型选择：WMCSFB

这是一个很长的故事……

• 我从 LeopardGecko 的自定义版本开始，这是用于体积分割的，使用 2D unets 和多平面方法。这是基于另一个名为 volume-segmantics 的类似包。效果非常差。
• wollny 3D unet。结果非常差。
• MedNext。 https://github.com/MIC-DKFZ/MedNeXt, https://arxiv.org/abs/2303.09975。结果非常好，所以我有一段时间定居在这个模型上，并用来改进训练算法。
• 随后使用 Wenzhao Zhao 的名为 WMCSFB 的 MedNext 类似模型进一步改进。参见参考文献 https://arxiv.org/abs/2402.16825，以及 GitHub https://github.com/ZhaoWenzhao/WMCSFB。这是基于 MedNext 的第一个版本，不支持 GRN，我修改了它以实现 MedNext v2 版本中的 GRN 功能。该模型的一个有趣特征是在 convnext/mednext 模型的第一层中使用球面傅里叶 - 贝塞尔函数（spherical Fourier-Bessel functions）进行初始化。模型代码全部包含在训练 notebook 中。

用于训练的代码提供在此处：https://www.kaggle.com/code/perdigao1/wmcsfb-grn-train?scriptVersionId=220076028

我使用的模型是小型 'S' 版本。更大版本的模型并没有带来更好的结果。

训练策略

• 未使用 copick 软件。我自己实现了掩码（masks）的创建。
• 掩码使用 1/3 的粒子半径生成。
• 最初使用 5 个体积进行训练。
• 我注意到的一个重要事项是，如果在训练期间使用体积 5_4 和 69_2，似乎会对最终分数产生不利影响，这表明这些数据或掩码存在问题。许多得分较高的 notebook 似乎也排除了这些体积用于训练。由于计算资源有限，我没有找到很好的解释。
• 训练体积：64x64x64 大小，粒子位于中心（除非是边缘粒子）。
• 向训练数据添加了随机体积，因此 2/3 的数据将是随机体积。
• 每个 64x64x64 体积简单归一化，使用均值和标准差。
• 增强：一个体素 shift，对比度 + 亮度，XY 平面 90 度旋转，沿所有轴翻转，体素交换（自定义编码），高斯噪声。
• 损失函数：自定义编码的 Tverski loss，支持类别权重 (0.25, 1.0, 1.0, 1.0, 4.0, 4.0)，分别对应 ( background, 'ribosome', 'virus-like-particle', 'apo-ferritin', 'beta-galactosidase', 'thyroglobulin' )（注意缺少"non-scoring"粒子）。
• metric, IoU (smp.utils.metrics.IoU())。
• 数据使用 sklearn 的 StratifiedKFold (4 折) 分割为训练集和验证集。
• 优化器：AdamW, Gradscaler。
• 使用 onecycleLR 调度器进行训练。
• epoch 数量和进一步结果呈现在 notebook 开头定义的 CFG 变量中。

随后进行了其他额外的训练调度，并且有一种机制可以通过 metric 分数选择“最佳”训练状态。然而，最好的结果是在这一步获得的，所以这里没有必要进一步阐述。

推理

https://www.kaggle.com/code/perdigao1/wmcsfb-grn-subm-inference

对于推理，体积被修补（blocks）为 (160,160,160) 的体积，重叠 (14,14,14)。除了边缘外，只考虑中间部分。结果（概率）然后进行 argmax，并运行连通分量（connected components）以获得质心。然后通过 DBSCAN 解析质心列表，这给出了略微更好的结果。DBSCAN 的想法来自一些分享的 notebook。

我获得的结果在视觉上相当不错。对于一些粒子，预测的质量似乎比 ground truth 本身更好，就位于正确的中心而言，或者有时似乎缺少注释。因此，我觉得很难想出更好的方法来改进。

我开始探索但未完成的是带有复数（相位）和四元数的 3D Unets。文献中有一些例子。一些数学直觉告诉我，这些 Unets 应该 工作得更好，因为完全训练的模型可能会收敛到类似于傅里叶或贝塞尔核的东西，就像这里使用的解决方案一样。然而，需要对应用归一化、残差和激活的方式进行重大更改。我尝试了一些草案版本，但这些要么训练不好，要么结果很差，公共分数最高为 0.5，作为一个开始还不错。