如何在PyTorch中進行多任務學習

在PyTorch中進行多任務學習可以使用多任務損失函數來同時優化多個任務。一種常用的方法是使用多個損失函數，每個損失函數對應一個任務，然后將這些損失函數進行加權求和作為最終的損失函數。下面是一個簡單的示例代碼：

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

# 定義多任務損失函數
class MultiTaskLoss(nn.Module):
    def __init__(self, task_weights):
        super(MultiTaskLoss, self).__init__()
        self.task_weights = task_weights
        
    def forward(self, outputs, targets):
        loss = 0
        for i in range(len(outputs)):
            loss += self.task_weights[i] * nn.CrossEntropyLoss()(outputs[i], targets[i])
        
        return loss

# 定義模型
class MultiTaskModel(nn.Module):
    def __init__(self):
        super(MultiTaskModel, self).__init__()
        self.fc1 = nn.Linear(10, 5)
        self.fc2 = nn.Linear(5, 2)
        
    def forward(self, x):
        x = self.fc1(x)
        output1 = self.fc2(x)
        output2 = self.fc2(x)
        
        return [output1, output2]

# 定義數據和標簽
data = torch.randn(1, 10)
target1 = torch.LongTensor([0])
target2 = torch.LongTensor([1])

# 創建模型和優化器
model = MultiTaskModel()
criterion = MultiTaskLoss([0.5, 0.5]) # 兩個任務的損失函數權重均為0.5
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 訓練模型
optimizer.zero_grad()
outputs = model(data)
loss = criterion(outputs, [target1, target2])
loss.backward()
optimizer.step()

在上面的示例中，我們定義了一個包含兩個任務的多任務模型和對應的多任務損失函數，其中兩個任務的損失函數權重均為0.5。在訓練過程中，我們計算模型輸出和目標值之間的損失，并根據總損失來更新模型參數。通過這種方式，我們可以實現多任務學習。