自动加权损失

这是 Liebel L、Körner M 的论文《多任务学习中的辅助任务》[J]（arXiv 预印本 arXiv:1805.06334，2018 年）的 PyTorch 实现。

上述论文改进了《利用不确定性对场景几何与语义的损失进行加权的多任务学习》一文（http://openaccess.thecvf.com/content_cvpr_2018/html/Kendall_Multi-Task_Learning_Using_CVPR_2018_paper.html），以避免训练过程中损失值变为负数。

需求

• Python
• PyTorch

如何使用您的模型进行训练

• 克隆仓库

git clone git@github.com:Mikoto10032/AutomaticWeightedLoss.git

• 创建 AutomaticWeightedLoss 模块

from AutomaticWeightedLoss import AutomaticWeightedLoss

awl = AutomaticWeightedLoss(2)  # 我们有 2 个损失
loss1 = 1
loss2 = 2
loss_sum = awl(loss1, loss2)

• 创建优化器来学习权重系数

from torch import optim

model = Model()
optimizer = optim.Adam([
                {'params': model.parameters()},
                {'params': awl.parameters(), 'weight_decay': 0}    
            ])

• 完整示例

from torch import optim
from AutomaticWeightedLoss import AutomaticWeightedLoss

model = Model()

awl = AutomaticWeightedLoss(2)  # 我们有 2 个损失
loss_1 = ...
loss_2 = ...

# 可学习参数
optimizer = optim.Adam([
                {'params': model.parameters()},
                {'params': awl.parameters(), 'weight_decay': 0}
            ])

for i in range(epoch):
    for data, label1, label2 in data_loader:
        # 前向传播
        pred1, pred2 = Model(data)    
        # 计算损失
        loss1 = loss_1(pred1, label1)
        loss2 = loss_2(pred2, label2)
        # 对损失进行加权
        loss_sum = awl(loss1, loss2)
        # 反向传播
        optimizer.zero_grad()
        loss_sum.backward()
        optimizer.step()

补充说明

实际上，这种方法并不总是有效，但我希望它能对您有所帮助。

AutomaticWeightedLoss 快速上手指南

AutomaticWeightedLoss 是一个基于 PyTorch 的开源工具，用于在多任务学习中自动调整不同损失函数的权重。它实现了论文《Auxiliary tasks in multi-task learning》中的算法，能够避免训练过程中损失值变为负数的问题，帮助模型更平衡地学习多个任务。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux / macOS / Windows
• 语言版本：Python 3.6+
• 核心依赖：PyTorch

请确保已安装 PyTorch。如果您在中国大陆地区，推荐使用清华源或中科大源加速安装：

pip install torch -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

本项目无需通过 pip 安装，直接克隆仓库即可使用。

1. 克隆项目仓库：

git clone git@github.com:Mikoto10032/AutomaticWeightedLoss.git

2. 将克隆后的目录添加到您的项目路径中，或者直接将 AutomaticWeightedLoss.py 文件复制到您的项目目录下。

基本使用

以下是集成该模块到训练流程的最简示例。假设您有两个需要同时优化的损失函数（loss1 和 loss2）。

1. 初始化模块与优化器

关键点在于需要将 awl 的参数也加入优化器中，并设置其 weight_decay 为 0。

from torch import optim
from AutomaticWeightedLoss import AutomaticWeightedLoss

# 初始化模型
model = Model()

# 初始化自动加权损失模块，参数 2 表示有 2 个损失函数
awl = AutomaticWeightedLoss(2)

# 构建优化器：同时优化模型参数和损失权重参数
optimizer = optim.Adam([
                {'params': model.parameters()},
                {'params': awl.parameters(), 'weight_decay': 0}	
            ])

2. 训练循环集成

在训练循环中，计算各个任务的原始损失后，传入 awl 对象即可获得加权后的总损失。

for i in range(epoch):
    for data, label1, label2 in data_loader:
        # 前向传播
        pred1, pred2 = Model(data)	
        
        # 计算原始损失
        loss1 = loss_1(pred1, label1)
        loss2 = loss_2(pred2, label2)
        
        # 自动加权求和
        loss_sum = awl(loss1, loss2)
        
        # 反向传播与更新
        optimizer.zero_grad()
        loss_sum.backward()
        optimizer.step()

通过以上步骤，模型将在训练过程中自动学习每个任务的最佳损失权重，无需手动调节超参数。