如何在TensorFlow Keras中实现指定双深度学习模型的训练与测试？

核心逻辑 #

训练阶段分为两个独立优化的模型：

• 蓝色模型（ModelA）：负责生成特征，基于自身任务损失更新参数。
• 红色模型（ModelB，黄色框部分）：以ModelA的输出作为输入，基于自身任务损失更新参数，且梯度不会反向传播至ModelA。
  测试阶段仅使用ModelB处理样本，无需加载ModelA。

代码示例（PyTorch） #

1. 模型定义 #

import torch
import torch.nn as nn
import torch.optim as optim

# 蓝色模型：特征提取+自身任务输出
class ModelA(nn.Module):
    def __init__(self, input_dim, feature_dim, task_a_output_dim):
        super().__init__()
        # 特征提取分支（输出给ModelB）
        self.feature_extractor = nn.Sequential(
            nn.Linear(input_dim, 128),
            nn.ReLU(),
            nn.Linear(128, feature_dim)
        )
        # 自身任务分支（用于计算损失）
        self.task_head = nn.Linear(feature_dim, task_a_output_dim)
    
    def forward(self, x):
        feature = self.feature_extractor(x)
        task_output = self.task_head(feature)
        return feature, task_output

# 红色模型：基于ModelA的特征完成自身任务
class ModelB(nn.Module):
    def __init__(self, feature_dim, task_b_output_dim):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(feature_dim, 64),
            nn.ReLU(),
            nn.Linear(64, task_b_output_dim)
        )
    
    def forward(self, x):
        return self.layers(x)

2. 训练流程 #

# 超参数设置
input_dim = 10          # 原始输入维度
feature_dim = 32        # ModelA输出的特征维度
task_a_dim = 5          # ModelA的任务输出维度（比如分类类别数）
task_b_dim = 2          # ModelB的任务输出维度（比如回归维度）

# 初始化模型、损失、优化器
model_a = ModelA(input_dim, feature_dim, task_a_dim)
model_b = ModelB(feature_dim, task_b_dim)

criterion_a = nn.CrossEntropyLoss()  # ModelA的损失（分类任务示例）
criterion_b = nn.MSELoss()           # ModelB的损失（回归任务示例）
optimizer_a = optim.Adam(model_a.parameters(), lr=1e-3)
optimizer_b = optim.Adam(model_b.parameters(), lr=1e-3)

# 训练循环
for epoch in range(100):
    # 模拟训练数据
    batch_x = torch.randn(32, input_dim)
    label_a = torch.randint(0, task_a_dim, (32,))  # ModelA的任务标签
    label_b = torch.randn(32, task_b_dim)          # ModelB的任务标签

    # 训练ModelA
    optimizer_a.zero_grad()
    feature, output_a = model_a(batch_x)
    loss_a = criterion_a(output_a, label_a)
    loss_a.backward()
    optimizer_a.step()

    # 训练ModelB：截断梯度，避免影响ModelA
    optimizer_b.zero_grad()
    # detach()断开特征与ModelA的计算图连接，阻止梯度回传
    feature_detached = feature.detach()
    output_b = model_b(feature_detached)
    loss_b = criterion_b(output_b, label_b)
    loss_b.backward()
    optimizer_b.step()

    # 打印训练状态
    if (epoch + 1) % 10 == 0:
        print(f"Epoch {epoch+1} | Loss A: {loss_a.item():.4f} | Loss B: {loss_b.item():.4f}")

3. 测试阶段 #

# 仅加载ModelB（无需ModelA）
model_b.eval()

with torch.no_grad():
    # 注意：测试输入需与ModelB训练时的输入格式一致
    # 场景1：若测试时用原始样本，需提前调整ModelB的输入维度为input_dim（修改ModelB的第一层Linear）
    test_x = torch.randn(16, input_dim)
    # 场景2：若测试时仍用ModelA生成的特征，需提前离线生成并保存，再输入给ModelB
    # test_feature = torch.randn(16, feature_dim)
    
    test_output = model_b(test_x)
    # 后续处理测试输出（如预测、评估等）

关键注意事项 #

• 梯度隔离：必须使用detach()或with torch.no_grad()处理ModelA的输出，确保ModelB的梯度不会反向传播到ModelA，这是实现两个模型独立训练的核心。
• 任务适配：根据实际任务选择对应的损失函数，比如分类用CrossEntropyLoss，回归用MSELoss。
• 测试输入一致性：如果训练时ModelB依赖ModelA的特征，而测试时不能使用ModelA，需提前调整ModelB的结构，让它可以直接处理原始输入（比如修改输入层维度，或训练时同时输入原始样本和ModelA特征进行融合学习）。
• 参数保存：训练完成后，分别保存ModelA和ModelB的参数，测试阶段仅加载ModelB的参数即可。

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