0
$\begingroup$

This is my basic multitask learning model, and it has 2 tasks. Since there are only 2 tasks, maybe I can duplicate the code for each task as self.tower1 and self.tower2.

class Multitask_Network(nn.Module):
    def __init__(self):
        super(Multitask_Network, self).__init__()

        in_ch = 20
        shared_layer_size = 10
        tower_hidden_size = 10
        output_size = 2

        self.sharedlayer = nn.Sequential(
            nn.Linear(in_ch, shared_layer_size),
            nn.ReLU(),
            nn.Dropout()
        )
        self.tower1 = nn.Sequential(
            nn.Linear(shared_layer_size, tower_hidden_size),
            nn.ReLU(),
            nn.Dropout(),
            nn.Linear(tower_hidden_size, output_size)
        )
        self.tower2 = nn.Sequential(
            nn.Linear(shared_layer_size, tower_hidden_size),
            nn.ReLU(),
            nn.Dropout(),
            nn.Linear(tower_hidden_size, output_size)
        )        

    def forward(self, x):
        x = self.sharedlayer(x)
        out1 = self.tower1(x)
        out2 = self.tower2(x)
        return out1, out2

model = Multitask_Network()
x = torch.ones((4, 20))
y = model(x)
y

But for my use case, I need to define the model's number of tasks during run time. I can do it with for-loop, but I think it is inefficient to run on GPU. I think it does the forward pass for task #1, then followed by task #2, to task #n, one after another.

How can I achieve this without a for loop in the forward pass?

class Multitask_Network(nn.Module):
    def __init__(self, num_tasks):
        super(Multitask_Network, self).__init__()

        in_ch = 20
        shared_layer_size = 10
        tower_hidden_size = 5
        output_size = 2
        self.num_tasks = num_tasks

        self.sharedlayer = nn.Sequential(
            nn.Linear(in_ch, shared_layer_size),
            nn.ReLU(),
            nn.Dropout()
        )

        self.conv_towers = []
        for i in range(num_tasks):

            tower = nn.Sequential(
                nn.Linear(shared_layer_size, tower_hidden_size),
                nn.ReLU(),
                nn.Dropout(),
                nn.Linear(tower_hidden_size, output_size)
            )

            self.conv_towers.append(
                    tower
                )
        self.conv_towers = nn.ModuleList(self.conv_towers)

    def forward(self, x):
        x = self.sharedlayer(x)

        output = []
        for i in range(self.num_tasks):
            output.append(self.conv_towers[i](x))

        return output

model = Multitask_Network(num_tasks=2)
x = torch.ones((4, 20))
y = model(x)
y
$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.