PyTorch如何正确优化双损失函数：马里奥赛车7自动驾驶模型训练慢问题

存在的严重问题 #

1. 模型结构设计不合理 #

现有结构将ResNet18的最终输出压缩到仅1维特征，再分别映射到加速和转向两个任务输出，1维特征完全不足以同时承载两个任务需要的视觉信息，两个任务会出现特征竞争，导致收敛极慢且效果上限极低。
修改方案：将ResNet18的全连接层输出维度调整为至少256维，再分别接两个任务的输出头，示例修改如下：

class MK7AutoPilot(nn.Module):
    def __init__(self):
        super().__init__()
        self.backbone = resnet18(pretrained=True)
        num_ftrs = self.backbone.fc.in_features
        # 输出256维通用特征
        self.backbone.fc = nn.Linear(num_ftrs, 256)
        # 加速分支：256->1->sigmoid
        self.accelerate_fc = nn.Linear(256,1)
        self.accelerate = nn.Sigmoid()
        # 转向分支：256->1->tanh
        self.steering_fc = nn.Linear(256,1)
        self.steering = nn.Tanh()

    def forward(self, inputs):
        x = self.backbone(inputs)
        x_accelerate = self.accelerate_fc(x)
        x_steering = self.steering_fc(x)
        accelerate =  self.accelerate(x_accelerate)
        steering = self.steering(x_steering)
        return accelerate, steering

2. 转向任务损失函数使用错误 #

HingeEmbeddingLoss是用于分类任务的损失函数，适合判断样本是否属于某一类、或者两个向量是否相似的场景，完全不适用于-1~1区间的连续值回归任务。
修改方案：转向损失替换为MSELoss或者SmoothL1Loss：

steering_loss = nn.MSELoss() # 或者 nn.SmoothL1Loss()

3. 训练代码存在多处关键bug #

• GPU调用无效：PyTorch中tensor.to(device)不会修改原张量，只会返回新的设备上的张量，现有写法所有数据都停留在CPU，没有用到GPU加速，训练速度自然极慢。修改为：

batch_images = batch_images.to(device)
acc_targets = acc_targets.to(device)
steer_targets = steer_targets.to(device)

• 验证集损失计算错误：running_loss仅在epoch开始时初始化一次，训练阶段累计的损失会被带入验证阶段计算，导致你看到的验证集损失远大于训练集，完全不符合真实情况。需要每个阶段开始前清零running_loss。
• 最优模型保存逻辑完全失效：lowest_loss初始值设为0.0，而你的实际损失都是几十到几百级别，永远不会触发更新最优模型的逻辑，最终保存的是epoch开始时的随机权重。需要将lowest_loss初始值改为无穷大float('inf')，且best_model_wts不需要每个epoch开始都保存，仅在验证损失更低时更新即可。
• 时间统计逻辑错误：time_elapsed的计算放在epoch循环外，仅统计最后一个epoch的耗时，要统计总训练时长需要把since = time.time()放在epoch循环外。

优化建议 #

• 微调预训练ResNet时建议初始学习率调整为1e-4~5e-5，原0.001的学习率过高，容易破坏预训练的特征权重。
• 可以先冻结ResNet18的 backbone 权重，仅训练最后的全连接头3~5个epoch，再解冻全网络训练，收敛速度会快很多。
• 可以给两个损失设置不同权重，比如loss = 0.3*acc_loss + 0.7*steer_loss，转向控制对自动驾驶更重要，可以适当提高其损失权重。

修正以上问题后不需要刻意延长训练时间，模型收敛速度会有数量级的提升，该方案本身是可行的，同类的自动驾驶游戏AI已有很多落地实现。

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