前言

本文介绍多模态融合中，如何解决传感器故障问题；基于激光雷达和相机，融合为BEV特征，实现3D检测和BEV分割，提高系统容错性和稳定性。

会讲解论文整体思路、模型框架、论文核心点、损失函数、实验与测试效果等。

论文地址：MetaBEV: Solving Sensor Failures for BEV Detection and Map Segmentation

代码地址：https://github.com/ChongjianGE/MetaBEV

一、模型框架

MetaBEV整体流程如下图所示：

模型的框架流程：

1. 输入数据，包括多视角相机的图像数据，激光雷达的点云数据。
2. 对图像数据进行特征提取，预测每个像素的深度，投影到三维视锥空间，然后经过BEV池化，得到图像BEV特征。
3. 对点云数据进行特征提取，使用体素化和稀疏3D卷积来编码，然后压缩Z维度，得到点云BEV特征。
4. 通过跨模态可变形注意力，融合图像BEV特征和点云BEV特征。
5. 经过几层的跨模态可变形注意力，然后来到自注意力层，进一步融合两种模态的信息。
6. 得到BEV融合特征，后面接两个任务头，包括3D物体检测头、BEV语义分割头。
7. 输出3D物体检测信息、BEV语义分割信息。

注意：这里的跨模态可变形注意力融合了MoE机制的，自注意力层也融合了MoE，后面详细讲。

BEV-Evolving Decoder，是论文的核心创新，由三个关键组成部分组成：跨模态注意力层、自注意力层和即插即用的M²oE块。

它用跨模态变形注意力机制，在相机和激光雷达间进行特征融合。然后经过自注意力层，进一步融合两种模态的信息。  

 二、核心内容——跨模态注意力+MoE

• 此层首先初始化一组密集的BEV查询，称为Bm​，每个查询代表特定空间网格内的特征。
• 这些查询通过位置编码（PE）进行增强，然后与相机BEV特征Bc​，或激光雷达BEV特征Bl​，或两者进行关联。
• 跨模态融合，采用了可变形注意力（Deformable Attention, DAttn(·)）。

2.1 跨模态注意力用公式表示：

• 可变形注意力机制，首先从BEV查询特征Bm​中生成偏移量Δpx和注意力权重Ax，然后结合Bm​的像素坐标和相应的偏移量来定位采样的值特征。
• 通过注意力权重Ax重新缩放采样特征后，Bm从传感器的信息特征中更新。
• 公式中，其中m表示注意力头，K表示采样的键的数量，p表示参考点。Wm​和Wm′​是相关的权重矩阵。
• 在公式中，输入特征x可以是Bc​、Bl​或两者的组合，这使MetaBEV能够灵活地在训练和测试期间使用多模态特征进行可变形注意力计算。

2.2 注意力中特定的MLP层：

• 原始的DAttn(·)实现不适合处理任意输入，因为它使用统一的MLP层来采样参考点Δpx和注意力权重Ax。
• 为了增加灵活性，引入了模型特定的MLP层（用于相机的C-MLP和用于激光雷达的L-MLP）
• 这里的"C-MLP"和"L-MLP"分别代表针对相机数据和激光雷达数据定制的MLP层。

通常，在一个多模态学习系统中，如果使用统一的MLP处理不同模态的数据，可能无法捕捉到每种模态独特的特征和数据分布。

例如，相机图像中的信息可能与颜色和纹理有关，而激光雷达数据可能更多地与空间和几何结构有关。

这两种类型的数据在特征空间中的表示可能完全不同。

如果使用同一个MLP层来处理这两种截然不同的数据类型，可能会导致模型无法充分学习到每种模态中的独特和有价值的信息。

为了解决这个问题，MetaBEV引入了模型特定的MLP层。这意味着：

• C-MLP（Camera-specific MLP）：专门设计用于处理相机BEV特征的MLP，它可以专注于学习和优化相机数据的处理，从而更好地捕捉图像中的色彩、纹理等视觉特征。

• L-MLP（LiDAR-specific MLP）：专门设计用于处理激光雷达BEV特征的MLP，它可以专注于学习和优化激光雷达数据的处理，从而更好地捕捉空间、几何等结构信息。

3.3 跨模态注意力集成MoE

如下图所示，是BEV-Evolving的结构图，展示了两种不同的Mixture of Experts (MoE)选项，并将这些结构集成到跨模态可变形注意力层中。

通过整合两种不同的MoE选项，可以在多任务学习中减少任务间的冲突。

RM²oE提供了一种通过路由机制选择专家的方式，而HM²oE提供了一种固定分配专家的方式。

首先看看RM²oE: Routing Multi-task Mixture-of-Experts

• 由不同的“专家”组成，每个专家负责特定的任务（检测或分割）。

• “路由器”负责决定对于每个任务应该使用哪个“专家”的输出。

• 结构内的“加权求和”(Weighted Sum)和“加法”(Add)操作用于合并来自不同专家的信息。

再看看HM²oE: Hard MoE

• 也是一个基于混合专家MoE的系统，但这里的每个专家是针对检测或分割任务分别设计的专用网络层（例如Det-FFN和Seg-FFN）

最后看看注意力中特定的MLP层

• 使用模态特定的MLP层来灵活地计算采样偏移和注意力权重。
• 引入了模态特定的注意力权重（如L-MLP用于激光雷达数据，C-MLP用于相机数据）
• 通过引入模态特定的MLP，这种结构允许系统灵活地处理不同类型的传感器数据，以及不同模态下的采样偏移和注意力权重计算。

三、核心内容——自注意力+MoE

自注意力层，它是BEV-Evolving解码器的一部分。

• 模型内部相关性：为了有效地建模查询之间的内部相关性，用BEV查询Bm​替换了之前公式中的输入特征 x。这意味着自注意力计算现在是基于BEV查询来进行的，从而能够捕捉不同查询（或特征）之间的相互关系。

• 自注意力层数：在BEV-Evolving解码器中，只组装了2 层自注意力层，这样做是为了在性能和计算效率之间取得一个理想的平衡。

• 模态间和查询内的关系建模：通过建模模态间（如相机和激光雷达）和查询内的相关性，融合的BEV特征最终被输出用于3D预测。

• 混合设计：实验发现，将内部和跨模态相关性结合起来的混合设计为融合的BEV特征提供了全面的建模，这对于多种任务都是有益的。

四、MoE混合专家

Mixture of Experts (MoE) 是一种机器学习架构，通过将多个专家网络（即专家）的预测结合起来，来处理复杂的任务。

这种方法的核心思想是：不同的专家可以学习并成为任务的不同子部分的专家，而一个门控网络负责决定在给定的输入下哪个专家的建议最应该被采纳。

MoE架构在多任务学习、大规模语言模型以及其他需要处理广泛数据分布的应用中特别有用。

MoE的工作原理：

• 专家网络：在MoE模型中，每个“专家”都是一个小型的神经网络，专门处理输入数据的一部分或某个特定任务。专家的设计允许模型在不同区域的输入数据上专业化，从而提高整体性能。
• 门控机制：门控网络负责学习如何根据输入数据分配给每个专家的权重。这意味着它决定了哪些专家对当前任务最为重要，并相应地调整它们的输出贡献。
• 输出融合：最后根据门控网络分配的权重，将选中的专家的输出加权融合，以产生最终预测。

MoE的优点：

• 灵活性：MoE可以轻松扩展以包括更多的专家和更复杂的门控机制，使其能够处理非常复杂和多样化的数据集。

• 专业化：通过将不同的专家专门用于数据的不同子集或不同任务，MoE能够在特定区域内达到更高的精确度。

• 高效性：尽管MoE包含多个专家，但由于只有一部分专家被激活处理特定的输入，所以它可以在不显著增加计算成本的情况下提供定制化的输出。

五、背景与设计动机

• 自动驾驶车辆的感知系统依赖于多模态传感器，如激光雷达和相机，以获得对周围环境的全面理解。这些传感器的信息被综合利用，以提高定位、导航和避障的准确性。
• 然而，实际应用中的传感器故障和数据损坏问题，会使当前方法存在局限性。

1）特征错位和CNN的限制

• 特征错位问题主要在多模态融合过程中，不同传感器数据的地理和几何对应关系难以精确匹配。
• 当前流行的基于卷积神经网络（CNN）的融合方法虽然在特征提取方面表现出色，但它们在处理长距离依赖和动态调整对不同特征的关注程度方面存在局限。
• 这是因为CNN主要关注局部特征，而自动驾驶场景下的物体识别和环境理解往往需要对更广泛区域的理解和整体上下文的分析。

2）完整模态的依赖性

• 多模态融合策略对完整模态输入的依赖表明，一旦特定传感器数据不可用或损坏，系统的性能会显著下降。
• 在现实世界中，传感器受环境因素影响而出现故障是不可避免的。例如，激光雷达可能因为雨雪或灰尘而失效，相机在低光照或直射阳光下可能无法捕捉清晰图像。
• 因此，依赖于特定模态的融合策略在实际应用中面临重大挑战。

MetaBEV认为模态依赖方法的主要瓶颈是缺乏设计，提出了一个模态任意的BEV-Evolving解码器，使得融合模块能够独立融合不同模态。

原理：使用跨模态注意力，将可学习的BEV查询与单个相机BEV特征、激光雷达BEV特征或两者相结合，以消除瓶颈。

总结：为了解决上述特征错位和完整模态依赖问题，MetaBEV通过在统一的BEV表示空间中，进行模态任意和任务不可知的学习。 

3）多任务学习性能冲突

• 考虑到实际中有限的计算资源，使用一个具有共享参数的单一框架来处理不同任务，比使用多个框架处理多个任务更为高效。
• 然而，检测和分割的联合学习中的任务冲突经常导致严重的性能下降，现有方法很少分析和设计针对多任务学习（MTL）。

受到专为自然语言处理和自监督领域中的大型语言模型设计的专家混合（MoE）的启发MetaBEV引入了一个具有新的M2oE-FFN层的鲁棒融合模块。

主要目的是为了减轻检测和分割之间的梯度冲突，以实现更加平衡的性能。

MetaBEV是第一个将MoE引入3D目标检测和BEV地图分割作为一种多模态、多任务和鲁棒方法的框架。

六、模型细节信息

1）模型结构设计

• 使用Swin-T和VoxelNet作为相机和激光雷达的特征编码器。
• BEV-Evolving解码器包含四个跨模态注意力层和两个自注意力层，用于生成融合的BEV。
• 初始化分辨率为180x180的meta-BEV以捕捉细粒度的关联。
• 使用FPN层生成融合BEV的多尺度特征。
• 默认情况下使用变换器头部（Transformer head）进行3D检测，CNN头部用于地图分割。

2）数据集和评估指标：

• 在nuScenes数据集上评估MetaBEV，这是一个大型多模态数据集。
• 数据集被分为训练/验证/测试三部分，包含多种传感器数据。
• 评估指标包括标准nuScenes检测得分（NDS）和平均精度（mAP）。
• BEV地图分割使用mIoU（平均交并比）评估。

3）训练配置：

• 遵循MMDetection3D的图像和激光雷达数据增强策略。
• 使用AdamW优化器，权重衰减为0.05，并使用循环学习率计划。
• 3D检测任务训练26个时代，地图分割任务训练20个时代。
• 使用CBGS策略平衡数据采样。
• MetaBEV在8个A100 GPUs上训练。
• 切换模态训练的比例设置为相机/激光雷达/两者的1/3。

4) 切换模态训练（Switched Modality Training）

这种训练策略提供了对不同传感器数据类型的灵活性，从而允许模型更好地适应可能发生的传感器失效或数据损坏情况。

• 模态特定模块: MetaBEV通过其BEV-Evolving块中的模态特定模块，可以灵活处理相机特征、激光雷达特征或两者的组合。

• 训练方案: 提出了一种切换模态的训练方案，这种方案可以使用任意模态进行精确预测。这种训练方法通过在训练过程中随机使用预先确定概率的不同模态输入来模拟真实世界条件。

• 任意模态推断: 作为结果，MetaBEV可以在任意输入模态上进行推断，这增加了它在自动驾驶中的实用性。

• 部署效率: 重要的是，这种方法只需要一组预先训练好的权重，就可以用于模型的所有部署场景。

路由机制在切换模态训练中起到了关键作用。

• 路由机制负责决定在给定输入时哪个模态特定的处理单元（或“专家”）应当被激活。
• 这个决策过程对于模型能否有效处理多模态数据至关重要。
• 在Mixture of Experts (MoE)模型中，路由器的设计至关重要，因为它决定了数据流向哪个专家，从而影响最终的输出质量和模型的鲁棒性。

5）传感器故障类型：

1. Limited Field (LF): 指的是由于硬件故障导致激光雷达（LiDAR）点云数据在某些角度上缺失，这种情况下数据收集不会覆盖完整的360度视角。在这项工作中，模拟了不同角度下的缺失LiDAR数据，包括360, 240, 180, 120度的情况。

2. Missing of Objects (MO): 在现实世界中，由于对象颜色和表面材质的不同等因素，激光雷达系统检测到的反射点数量可能会减少。这项工作通过按特定比例概率移除对象的点云数据来模拟对象点缺失，选定的比例包括0.0, 0.1, 0.3, 0.5, 0.7, 1.0。

3. Beam Reduction (BR): 感测系统中发生的一种现象，当可用的电源供应或传感器处理能力受限时，系统无法处理传感器收集的全部数据，导致处理的激光束数量减少。为了评估，选择了1, 4, 8, 16, 32束激光束进行评估。

4. View Noise (VN): 由于传感器噪声、电气干扰、大气条件或压缩伪影等多种因素，捕获的视图中包含随机变化或扭曲，这些变化或扭曲并不代表实际场景。通过用随机生成的噪声替换0到6个视图来模拟视觉噪声。

5. View Drop (VD): 指的是现实世界中的一种情况，摄像头没有捕捉到部分视觉场景，导致信息丢失。这可能是由于摄像头定位不正确或硬件故障等多种因素造成的。以类似于VN的方式，通过使用零初始化的输入来替代多达六个缺失的视图进行模拟。

6. Obstacle Occlusion (OO): 当场景中的对象部分被遮挡物或遮挡所遮盖时发生的感知现象。通过使用一组预定义的蒙版进行alpha混合来模拟这种现象，从而生成一个被遮挡的视图。

七、实验测试与效果

7.1）在nuScenes验证集上的性能比较，包括3D检测、分割。

• Modality: 表示使用的传感器类型，C表示相机，L表示激光雷达，L+C表示两者的组合。
• MTL: 表示是否使用了多任务学习（Multi-Task Learning），勾选表示是。
• mAP(val): 表示在验证集上的平均精度（mean Average Precision），是3D检测任务的一个常用指标。
• NDS(val): NuScenes检测评分（NuScenes Detection Score），是NuScenes数据集专用的综合性能评价指标。
• Drivable/Ped.Cross/Walkway 等: 表示在特定类别上的分割性能，比如可驾驶区域、行人穿越道、人行道等。
• Mean: 所有场景分割性能的平均值。

7.2）MetaBEV在面对极端传感器缺失情况下的性能表现，与其他方法（如TransFusion和BEVFusion）进行了比较。

• 这些评估基于三种不同的场景：正常使用相机和激光雷达的情况，缺失相机，以及缺失激光雷达。
• MetaBEV系统能够在传感器数据缺失时，完全不使用缺失的特征进行推断，表现出较强的鲁棒性。

7.3）自动驾驶的感知模型中，考虑不同类型的传感器故障：

• Limited Field of LiDAR (LF): 由于收集错误或硬件损坏，激光雷达数据只能从视场的一部分收集到。

• Missing of Objects (MO): 由于某些物质阻止激光雷达点被反射，导致无法检测到一些物体。

• Beams Reduction (BR): 由于供电或传感器处理能力有限，激光雷达的光束数量减少。

• View Drop (VD) 和 View Noise (VN): 这两种都是由于相机故障造成的，VD是指相机视图突然丢失，而VN是指相机视图带有噪声。

• Obstacle Occlusion (OO): 这是一个现实世界中的现象，即物体被遮挡导致相机无法检测到。

• Missing Camera(MC): 相机完全丢失的情况。

• Missing LiDAR(ML): 激光雷达完全丢失的情况。 

下面表格展示了MetaBEV和BEVFusion在各种传感器损坏情况下的性能比较 ：

• 在零样本测试中，MetaBEV在所有类型的传感器损坏情况下的NDS和mIoU性能均优于BEVFusion。
• 在领域内测试中，MetaBEV同样展示了优于BEVFusion的性能，尤其在mIoU指标上有显著提升。
• 在视野限制和物体遮挡的情况下，MetaBEV的性能提升尤为明显，这表明MetaBEV在处理这些特定损坏情况下具有更强的鲁棒性。

补充介绍：

• 零样本测试（zero-shot test）：这种测试是为了评估模型处理前所未见的损坏或缺失数据的能力。MetaBEV在没有见过的损坏数据上直接评估训练好的模型。
• 领域内测试（in-domain test）：在这种测试中，MetaBEV在随机损坏的数据上进行训练，然后在带有特定噪声等级的数据上进行评估。
• 在这两种测试中，MetaBEV在11/12种传感器损坏情形的零样本评估上都优于BEVFusion。

7.4）MetaBEV在多任务学习的实验和结果，比较了两种不同的专家混合（MoE）结构。

• 实验考虑了两种MoE结构：H-MoE和R-MoE。
• 在仅进行检测任务时，mAP为67.6%，NDS为71.0%。
• 在仅进行分割任务时，mIoU为70.4%。
• H-MoE在多任务学习环境中提供了一定程度的性能提升，R-MoE提供了更显著的性能增益。

7.5） MetaBEV模型的消融研究，以评估不同的架构设计和训练策略对性能的影响。

• 图(a)展示了不同数量的跨模态和自注意力层对mAP和NDS的影响。 结果显示，“2自注意力层+4跨模态注意力层”的组合在NDS指标上达到了最高值69.3。
• 图(b)测试了不同参考点数量对性能的影响。 结果表明，使用8个参考点时，NDS达到最高值68.2。
• 图(c)研究了不同数量的专家对mAP和mIoU的影响。 结果显示，“1/2 Expert”设置在mAP上达到了最高值69.8，而“1/8 Expert”设置在mIoU上达到了最高值66.4。
• 图(d)对比了使用标准训练（Vanilla training）与切换模态训练（Switched Modality Training）在传感器完全失效时的表现。
  • 在LiDAR完全失效的情况下， 使用切换模态训练的NDS和mAP指标显著提高（+32.8和+36.1）。
  • 在相机完全失效的情况下， 切换模态训练同样显著提高了性能（+1.0和+1.2）。
  • 在模态齐全的情况下（Full）， 切换模态训练在NDS上有所提高，但在mIoU上略有下降（-1.9）。

小结： 

• 发现几层跨模态注意力层足以产生有效的预测，而增加更多层并不一定能提高模型的能力。
• 发现仅增加自注意力层可以显著提高预测性能。
• 使用2/8专家配置（在八个专家中选择得分最高的两个）可以获得最佳性能。
• 使用切换模态训练策略，模拟传感器缺失的情况，与常规全模态训练相比，能够显著提高3D检测和BEV地图分割的性能。
• 在LiDAR缺失和相机缺失的场景下，分别有+32.8%和+1.0%的NDS改善，以及分割性能的显著提升。
• 甚至在完整模态下，这种训练策略也能从71.0%的NDS提升到71.5%。

补充一下：

• “Expert number”是在Mixture of Experts (MoE)模型中专家（expert）的数量。
• 在表格中的“Expert number”列可能展示了不同的配置，比如使用“1/2专家”、“1/4专家”、“1/8专家”等，这些配置指的是在模型中有多少比例的专家参与了最终的决策过程。
• 这里的比例通常是用于控制模型复杂度和计算资源分配的一个方式。

7.6）MetaBEV和BEVFusion 在不同摄像头噪声条件下的地图分割结果的比较

不同的环境条件包括：

• 亮度调整（Brightness）
• 对比度调整（Contrast）
• 饱和度调整（Saturation）
• JPEG压缩噪声（JPEG Compression）
• 雪（Snow）
• 雾（Fog）
• 霜冻（Frost）
• 飞溅（Splatter）
• 运动模糊（Motion Blur）
• 焦外模糊（Defocus Blur）
• 玻璃模糊（Glass Blur）
• 高斯模糊（Gaussian Blur）

• 使用的指标是mean Intersection over Union (mIoU)，它是衡量分割质量的常用指标。 
• 左侧雷达图展示了不同图像退化条件下的mIoU分数。
• 右侧雷达图展示了相同条件下模型性能保持的程度。

在左图中，MetaBEV在大多数条件下相比BEVFusion有更高的得分，这表明MetaBEV在处理图像噪声时能够更好地保持地图分割性能。

例如，在亮度和对比度调整条件下，MetaBEV的性能明显优于BEVFusion。

右图显示的性能保持度（retention）与左图类似，也显示了MetaBEV在大多数条件下的性能都优于BEVFusion。

模型效果：

八、关键模块代码

Cross-Attention Modules 示例代码如下所示：

import torch
import torch.nn as nn

class CrossAttention(nn.Module):
    def __init__(self, d_model=264, n_heads=8, n_points=4, lidar_flag=False, camera_flag=False, **kwargs):
        super().__init__()
        if camera_flag:
            self.offsets_camera = nn.Linear(d_model, n_heads * n_points * 2)
            self.attention_weights_camera = nn.Linear(d_model, n_heads * n_points)
        
        if lidar_flag:
            self.offsets_lidar = nn.Linear(d_model, n_heads * n_points * 2)
            self.attention_weights_lidar = nn.Linear(d_model, n_heads * n_points)
        
        self.value_proj = nn.Linear(d_model, d_model)
        self.output_proj = nn.Linear(d_model, d_model)
    
    def forward(self, query, reference_points, x, camera_flag=False, lidar_flag=False):
        value = self.value_proj(x)
        if camera_flag and lidar_flag:
            offsets_camera = self.offsets_camera(query)
            offsets_lidar = self.offsets_lidar(query)
            offsets = torch.cat([offsets_camera, offsets_lidar])
            attention_weights_camera = self.attention_weights_camera(query)
            attention_weights_lidar = self.attention_weights_lidar(query)
            attention_weights = torch.cat([attention_weights_camera, attention_weights_lidar], dim=3)
        elif camera_flag and not lidar_flag:
            offsets = self.offsets_camera(query)
            attention_weights = self.attention_weights_camera(query)
        elif lidar_flag and not camera_flag:
            offsets = self.offsets_lidar(query)
            attention_weights = self.attention_weights_lidar(query)
        
        attention_weights = F.softmax(attention_weights)
        sampling_locations = reference_points + offsets
        
        # MSDeformAttnFunction is implemented by DeformableAttention
        output = MSDeformAttnFunction.apply(value, sampling_locations.float(), attention_weights, output_mask)
        output = self.output_proj(output)
        return x

HMoE示例代码如下所示：

import torch
import torch.nn as nn
import torch.nn.functional as F

# HMoE: Hard MoE
class MultiTaskExpertMlp(nn.Module):
    def __init__(self, in_features=256, hidden_features=None, out_features=None, act_layer=nn.GELU(), drop=0., det_flag=False, seg_flag=False, **kwargs):
        super().__init__()
        self.det = det_flag
        self.seg = seg_flag
        self.multi_tasks = det_flag and seg_flag
        
        if self.multi_tasks:
            self.det_expert = nn.Sequential(
                nn.Linear(in_features, hidden_features), act_layer, nn.Dropout(drop),
                nn.Linear(hidden_features, out_features), nn.Dropout(drop)
            )
            self.seg_expert = nn.Sequential(
                nn.Linear(in_features, hidden_features), act_layer, nn.Dropout(drop),
                nn.Linear(hidden_features, out_features), nn.Dropout(drop)
            )
            self.fuse = nn.Sequential(nn.Linear(out_features*2, out_features),)
            self.norm = nn.LayerNorm(out_features)
        elif self.seg:
            self.seg_expert = nn.Sequential(
                nn.Linear(in_features, hidden_features), act_layer, nn.Dropout(drop),
                nn.Linear(hidden_features, out_features), nn.Dropout(drop)
            )
        elif self.det:
            self.det_expert = nn.Sequential(
                nn.Linear(in_features, hidden_features), act_layer, nn.Dropout(drop),
                nn.Linear(hidden_features, out_features), nn.Dropout(drop)
            )

    def forward(self, x):
        b, n, c = x.shape
        if self.multi_tasks:
            det_fea = self.det_expert(x)
            seg_fea = self.seg_expert(x)
            fuse_fea = self.norm(self.fuse(torch.cat([det_fea, seg_fea], dim=2)))
            return fuse_fea
        elif self.det:
            x = self.det_expert(x)
            return x
        elif self.seg:
            x = self.seg_expert(x)
            return x

RMoE示例代码如下所示：

import torch
import torch.nn as nn

# RMoE: Routing MoE
from tutel import moe as tutel_moe
from tutel import moe_utils

class MultiTaskRouteMoEMlp(nn.Module):
    def __init__(self, in_features=256, hidden_features=None, out_features=None, drop=0., act_layer=nn.GELU(), **kwargs):
        super().__init__()
        self.act = act_layer
        gate_type = kwargs['gate_type']
        experts = kwargs['experts']

        # Implemented using Microsoft-tutel: https://github.com/microsoft/tutel
        self._moe_layer = tutel_moe.moe_dp_layer(
            gate_type=gate_type,
            experts={
                'type': experts.type,
                'count_per_node': experts.count_per_node,
                'hidden_size_per_expert': hidden_features,
                'activation_fn': lambda x: self.act(x),
                'dropout_fn': lambda x: nn.Dropout(drop)(x)
            },
            model_dim=out_features,
            scan_expert_func=None,
            group=net.create_groups_from_world(group_count=1).data_group, # used for DistributedDataParallel mode
            seeds=(1, dist.get_rank() + 1, 1),
            a2a_ffn_overlap_degree=experts.a2a_ffn_overlap_degree,
            parallel_type=experts.parallel_type,
            use_2d=experts.use_2dh,
        )

    def forward(self, x, gate_index):
        return self._moe_layer(x, gate_index=gate_index)

本文先介绍到这里，后面会分享“多模态融合”的其它数据集、算法、代码、具体应用示例。