如何使用 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()

推荐答案

您可以(并且应该)使用 nn 容器，例如 <代码>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屋！