多任务学习 的代码

以下是一个使用PyTorch实现多任务学习的示例代码：

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, 3)
        self.task2_layer = nn.Linear(5, 2)

    def forward(self, x):
        x = self.shared_layer(x)
        task1_output = self.task1_layer(x)
        task2_output = self.task2_layer(x)
        return task1_output, task2_output

# 定义数据集
x1 = torch.randn(100, 10)
y1 = torch.randint(0, 3, (100,))
x2 = torch.randn(100, 10)
y2 = torch.randint(0, 2, (100,))
train_data = [(x1, y1), (x2, y2)]

# 定义损失函数和优化器
model = MultiTaskModel()
criterion1 = nn.CrossEntropyLoss()
criterion2 = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

# 训练多任务模型
for epoch in range(10):
    for x, y in train_data:
        optimizer.zero_grad()
        task1_output, task2_output = model(x)
        loss1 = criterion1(task1_output, y1)
        loss2 = criterion2(task2_output, y2)
        total_loss = loss1 + loss2
        total_loss.backward()
        optimizer.step()

在上面的代码中，我们定义了一个名为`MultiTaskModel`的多任务模型，该模型共享一层线性层（`shared_layer`），并针对不同的任务分别定义了不同的输出层（`task1_layer`和`task2_layer`）。我们还定义了两个损失函数（`criterion1`和`criterion2`）和一个优化器（`optimizer`），并使用这些组件训练模型。在每个训练迭代中，我们分别计算两个任务的损失，并将它们相加以获得总体损失。最后，我们使用反向传播来更新模型的参数。