如何使用pytorch构建多任务DNN，例如用于100多个任务? [英] How to use pytorch to construct multi-task DNN, e.g., for more than 100 tasks?

问题描述

下面是使用pytorch为两个回归任务构造DNN的示例代码. forward 函数返回两个输出(x1，x2).网络如何处理许多回归/分类任务?例如100或1000个输出.对所有输出(例如x1，x2，...，x100)进行硬编码绝对不是一个好主意.有没有简单的方法可以做到这一点?谢谢.

Below is the example code to use pytorch to construct DNN for two regression tasks. The forward function returns two outputs (x1, x2). How about the network for lots of regression/classification tasks? e.g., 100 or 1000 outputs. It definitely not a good idea to hardcode all the outputs (e.g., x1, x2, ..., x100). Is there an simple method to do that? Thank you.

import torch
from torch import nn
import torch.nn.functional as F

class mynet(nn.Module):
    def __init__(self):
        super(mynet, self).__init__()
        self.lin1 = nn.Linear(5, 10)
        self.lin2 = nn.Linear(10, 3)
        self.lin3 = nn.Linear(10, 4)

    def forward(self, x):
        x = self.lin1(x)
        x1 = self.lin2(x)
        x2 = self.lin3(x)
        return x1, x2

if __name__ == '__main__':
    x = torch.randn(1000, 5)
    y1 = torch.randn(1000, 3)
    y2 = torch.randn(1000,  4)
    model = mynet()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-4)
    for epoch in range(100):
        model.train()
        optimizer.zero_grad()
        out1, out2 = model(x)
        loss = 0.2 * F.mse_loss(out1, y1) + 0.8 * F.mse_loss(out2, y2)
        loss.backward()
        optimizer.step()

推荐答案

您可以(并且应该)使用code> nn.ModuleList 或 nn.ModuleDict 来管理任意数量的子模块.

You can (and should) use nn containers such as nn.ModuleList or nn.ModuleDict to manage arbitrary number of sub-modules.

例如(使用 nn.ModuleList >):

For example (using nn.ModuleList):

class MultiHeadNetwork(nn.Module):
    def __init__(self, list_with_number_of_outputs_of_each_head):
        super(MultiHeadNetwork, self).__init__()
        self.backbone = ...  # build the basic "backbone" on top of which all other heads come
        # all other "heads"
        self.heads = nn.ModuleList([])
        for nout in list_with_number_of_outputs_of_each_head:
            self.heads.append(nn.Sequential(
              nn.Linear(10, nout * 2),
              nn.ReLU(inplace=True),
              nn.Linear(nout * 2, nout)))

    def forward(self, x):
        common_features = self.backbone(x)  # compute the shared features
        outputs = []
        for head in self.heads:
            outputs.append(head(common_features))
        return outputs

请注意，在此示例中，每个磁头比单个 nn.Linear 层要复杂得多.
不同的头数"(和输出数)由参数 list_with_number_of_outputs_of_each_head 的长度确定.

Note that in this example each head is more complex than a single nn.Linear layer.
The number of different "heads" (and number of outputs) is determined by the length of the argument list_with_number_of_outputs_of_each_head.

重要通知:使用>> nn 容器，而不是简单的pythonic列表/字典来存储所有子模块.否则pytorch将难以管理所有子模块.
参见例如此答案，这个问题和这个问题.

这篇关于如何使用pytorch构建多任务DNN，例如用于100多个任务?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！