pytorch两个任务联合训练模型中损失函数权重的自动调整代码

下面是一个示例代码，展示如何通过自适应调整损失权重，来实现两个任务的联合训练：

import torch
from torch import nn, optim
from torch.utils.data import DataLoader

# 定义模型
class JointModel(nn.Module):
    def __init__(self, task1_model, task2_model):
        super(JointModel, self).__init__()
        self.task1_model = task1_model
        self.task2_model = task2_model

    def forward(self, input1, input2):
        output1 = self.task1_model(input1)
        output2 = self.task2_model(input2)
        return output1, output2

# 定义数据集
class JointDataset(torch.utils.data.Dataset):
    def __init__(self, task1_dataset, task2_dataset):
        self.task1_dataset = task1_dataset
        self.task2_dataset = task2_dataset

    def __getitem__(self, index):
        task1_data = self.task1_dataset[index]
        task2_data = self.task2_dataset[index]
        return task1_data[0], task2_data[0], task1_data[1], task2_data[1]

    def __len__(self):
        return len(self.task1_dataset)

# 定义训练函数
def train(model, dataloader, criterion1, criterion2, optimizer, device):
    model.train()
    total_loss = 0
    total_task1_loss = 0
    total_task2_loss = 0
    for input1, input2, target1, target2 in dataloader:
        input1 = input1.to(device)
        input2 = input2.to(device)
        target1 = target1.to(device)
        target2 = target2.to(device)

        # 前向传播
        output1, output2 = model(input1, input2)

        # 计算损失
        task1_loss = criterion1(output1, target1)
        task2_loss = criterion2(output2, target2)
        total_task1_loss += task1_loss.item()
        total_task2_loss += task2_loss.item()

        # 自适应调整损失权重
        loss_ratio = task1_loss.item() / task2_loss.item()
        criterion1.weight = torch.tensor([loss_ratio])
        criterion2.weight = torch.tensor([1.0 / loss_ratio])

        # 计算总损失
        loss = task1_loss + task2_loss
        total_loss += loss.item()

        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    return total_loss / len(dataloader), total_task1_loss / len(dataloader), total_task2_loss / len(dataloader)

# 定义数据集和模型
task1_dataset = ...
task2_dataset = ...
joint_dataset = JointDataset(task1_dataset, task2_dataset)

task1_model = ...
task2_model = ...
joint_model = JointModel(task1_model, task2_model)

# 定义损失函数和优化器
task1_criterion = nn.CrossEntropyLoss()
task2_criterion = nn.MSELoss()
optimizer = optim.SGD(joint_model.parameters(), lr=0.01)

# 训练模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
joint_model.to(device)
train_loader = DataLoader(joint_dataset, batch_size=32, shuffle=True)
for epoch in range(10):
    loss, task1_loss, task2_loss = train(joint_model, train_loader, task1_criterion, task2_criterion, optimizer, device)
    print("Epoch {}, Total Loss: {}, Task 1 Loss: {}, Task 2 Loss: {}".format(epoch+1, loss, task1_loss, task2_loss))

在上面的示例中，定义了一个 `JointModel` 类，用于组合两个任务的模型。定义了一个 `JointDataset` 类，用于组合两个任务的数据集。在训练函数中，使用 `criterion1.weight` 和 `criterion2.weight` 自适应调整两个任务的损失权重，以实现更好的联合训练效果。最后，通过 `train` 函数训练模型，并输出每个 epoch 的总损失和两个任务的损失。