多任务学习pytorch

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

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

# 定义模型
class MultiTaskModel(nn.Module):
    def __init__(self, num_classes_task1, num_classes_task2):
        super(MultiTaskModel, self).__init__()
        self.shared_layer = nn.Sequential(
            nn.Linear(input_shape, 64),
            nn.ReLU(),
            nn.Dropout(p=0.2)
        )
        self.task1_output = nn.Linear(64, num_classes_task1)
        self.task2_output = nn.Linear(64, num_classes_task2)
        
    def forward(self, x):
        x_shared = self.shared_layer(x)
        x_task1 = self.task1_output(x_shared)
        x_task2 = self.task2_output(x_shared)
        return x_task1, x_task2

# 实例化模型
model = MultiTaskModel(num_classes_task1, num_classes_task2)

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

# 训练模型
for epoch in range(num_epochs):
    for i, (inputs, labels_task1, labels_task2) in enumerate(train_loader):
        optimizer.zero_grad()
        outputs_task1, outputs_task2 = model(inputs)
        loss_task1 = criterion(outputs_task1, labels_task1)
        loss_task2 = criterion(outputs_task2, labels_task2)
        loss = loss_task1 + loss_task2
        loss.backward()
        optimizer.step()

# 评估模型
with torch.no_grad():
    for inputs, labels_task1, labels_task2 in test_loader:
        outputs_task1, outputs_task2 = model(inputs)
        _, predicted_task1 = torch.max(outputs_task1.data, 1)
        _, predicted_task2 = torch.max(outputs_task2.data, 1)
        total_task1 += labels_task1.size(0)
        correct_task1 += (predicted_task1 == labels_task1).sum().item()
        total_task2 += labels_task2.size(0)
        correct_task2 += (predicted_task2 == labels_task2).sum().item()
    acc_task1 = 100 * correct_task1 / total_task1
    acc_task2 = 100 * correct_task2 / total_task2

在上面的代码中，我们首先定义了一个多任务模型，其中共享一些层，同时每个任务都有自己的输出层。我们使用交叉熵损失函数和Adam优化器，训练和评估时分别传递每个任务的标签。在训练过程中，我们使用PyTorch的数据加载器来加载训练数据。