如何基于多条件构建带有多输出分支的PyTorch模型？

你当前的写法首先会遇到变量未定义报错：分支未触发时something_1/something_2/something_3不存在，根本无法走到分类器计算步骤。不需要调整返回值的数量，按以下方案修改即可：

• 第一步：提前初始化所有分支的中间变量
  所有something系列变量在进入条件判断前，先初始化和分支触发后shape完全一致的占位张量，推荐用和输入同设备的0张量，既避免变量未定义报错，也能保证后续分类器输入shape始终稳定。
• 第二步：新增有效性掩码同步返回
  返回值固定保留3路logits输出，额外新增3个对应维度的bool掩码，标记每个logits的对应位置是否有效。无效位置的logits是0输入计算得到的结果，后续计算损失时直接用掩码屏蔽即可，完全不会影响训练效果。
• 第三步：不要做动态返回值处理
  深度学习框架的静态图导出、批处理逻辑都不支持动态变化的返回值数量，且一个batch内可能同时存在满足不同条件的样本，固定返回结构不管是训练还是部署都会省事很多。

参考代码示例 #

修改后的forward逻辑参考：

def forward(self, x, condition1, condition2):
    batch_size = x.shape[0]
    # 提前初始化所有中间变量，shape和分支输出对齐
    something_1 = torch.zeros(batch_size, self.hidden_dim, device=x.device)
    something_2 = torch.zeros_like(something_1)
    something_3 = torch.zeros_like(something_1)

    if condition1 and condition2:
        something_1 = 分支1的计算逻辑
    if condition1:
        something_2 = 分支2的计算逻辑
    if condition2:
        something_3 = 分支3的计算逻辑

    o1 = self.classifier_output(something_1)
    o2 = self.classifier_output(something_2)
    o3 = self.classifier_output(something_3)

    # 生成对应掩码，批处理场景下可扩展为样本级掩码
    mask1 = (condition1 & condition2).view(-1, 1).expand_as(o1)
    mask2 = condition1.view(-1, 1).expand_as(o2)
    mask3 = condition2.view(-1, 1).expand_as(o3)

    return {
        "logits": [o1, o2, o3],
        "valid_masks": [mask1, mask2, mask3]
    }

后续损失计算只需要屏蔽无效位置即可：

loss_fn = nn.CrossEntropyLoss(reduction="none")
# 仅计算有效位置的损失
loss1 = (loss_fn(o1, y1) * mask1[:, 0]).mean()
loss2 = (loss_fn(o2, y2) * mask2[:, 0]).mean()
loss3 = (loss_fn(o3, y3) * mask3[:, 0]).mean()
total_loss = loss1 + loss2 + loss3

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