如何在PyTorch中实现多任务学习场景下的神经网络门控机制？

下面是两种兼顾效率、可维护性的主流实现方式：

• 方案1：批量预分组（训练效率最高，无冗余计算，适配任务数少的场景）
  训练前不需要修改单样本逻辑，只需要对每个batch的样本按任务激活标签（本场景为gender）做张量级拆分，全程用GPU并行运算，没有任何冗余计算：
  1. 先把整个batch的输入特征输入共享底层MLP，得到所有样本的共享表征
  2. 用布尔索引按gender拆分共享表征为男性子张量、女性子张量
  3. 仅把对应子张量输入对应性别的MLP分支，计算预测值和损失
  4. 两个分支的有效损失加总后统一反向传播即可
     代码层面实现非常简洁，以PyTorch为例：

  shared_out = shared_bottom(batch_x)
  # 生成样本分组掩码
  male_mask = batch_gender == 1
  female_mask = ~male_mask
  # 仅计算存在对应样本的分支损失
  male_loss = loss_fn(male_mlp(shared_out[male_mask]), male_label[male_mask]) if male_mask.any() else 0
  female_loss = loss_fn(female_mlp(shared_out[female_mask]), female_label[female_mask]) if female_mask.any() else 0
  total_loss = male_loss + female_loss
  total_loss.backward()
  

  该方案完全规避了无用分支的计算，也不需要写逐样本的if-else逻辑，GPU并行效率拉满，是双任务场景下的最优选择。
• 方案2：任务路由层封装（适配K>2的大规模多任务场景，扩展性更强）
  如果任务数量较多（比如大于10个），可以封装统一的任务路由层：提前将所有任务分支注册到路由层，输入共享表征和任务激活标签后，路由层自动按标签分组、调用对应分支计算。本质是对方案1的通用封装，新增任务时只需要注册对应分支即可，不需要修改计算逻辑，适合工业级大规模多任务系统。

补充说明你提到的两个原有方案的适用场景：掩码置零损失的方案适合样本分布极不均衡、没法稳定拆分子batch的场景，优点是代码改动最小，不需要调整数据加载逻辑；逐样本if-else的方案会严重破坏GPU并行性，仅适合小批量CPU训练，确实不推荐使用。

内容的提问来源于stack exchange，提问作者avocado