多任务学习框架中的损失效应

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

class MultiTaskModel(nn.Module):
    def __init__(self):
        super(MultiTaskModel, self).__init__()
        self.shared_layer = nn.Linear(10, 5)
        self.task1_layer = nn.Linear(5, 1)
        self.task2_layer = nn.Linear(5, 1)
    
    def forward(self, x):
        shared_output = torch.relu(self.shared_layer(x))
        task1_output = self.task1_layer(shared_output)
        task2_output = self.task2_layer(shared_output)
        return task1_output, task2_output

model = MultiTaskModel()
criterion_task1 = nn.MSELoss()
criterion_task2 = nn.BCEWithLogitsLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 假设我们有一些输入数据x和对应的标签y1, y2
x = torch.randn(32, 10)
y1 = torch.randn(32, 1)
y2 = torch.randint(0, 2, (32, 1)).float()

for epoch in range(10):
    optimizer.zero_grad()
    task1_output, task2_output = model(x)
    loss_task1 = criterion_task1(task1_output, y1)
    loss_task2 = criterion_task2(task2_output, y2)
    total_loss = loss_task1 + loss_task2
    total_loss.backward()
    optimizer.step()
    print(f'Epoch {epoch+1}, Loss Task1: {loss_task1.item()}, Loss Task2: {loss_task2.item()}')