Title: Multi-Task Learning and HydraNets with PyTorch - PyImageSearch

URL Source: https://pyimagesearch.com/2022/08/17/multi-task-learning-and-hydranets-with-pytorch/

Published Time: 2022-08-17T13:00:00+00:00

Markdown Content:
* * *

**Table of Contents**
---------------------

*   [Multi-Task Learning and HydraNets with PyTorch](#h2HydraNets)
*   [Solving a Multi-Task Learning Project](#h2Solving)

*   [Creating a Multi-Task DataLoader with PyTorch](#h3Creating)

*   [Data](#h4Data)
*   [Dataset](#h4Dataset)
    [](#h4Dataset)*   [](#h4Dataset)[INIT](#h5INIT)
    *   [LEN](#h5LEN)
    *   [GET\_ITEM](#h5GET_ITEM)
*   [DataLoaders](#h4DataLoaders)

*   [Model Definition](#h3Definition)

*   [Training: _Loss Functions_](#h4Functions)
*   [Training Loop](#h4Loop)

*   [Conclusion](#h2Conclusion)

*   [Citation Information](#h3Citation)

* * *

Today, we will learn about Multi-Task Learning and **HydraNets**. This is a Deep Learning technique I first introduced back in mid-2020 in an email I sent to exactly 653 people.

![Image 1](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20700%20393'%3E%3C/svg%3E)

**The responses to this email were so high** (engineers from everywhere around the planet told me they loved it and wanted to apply it to their company) **I had to create an entire HydraNet section in my course catalog**. You can learn more by visiting [https://www.thinkautonomous.ai/](https://www.thinkautonomous.ai/)

Not only is this technique **new** and **exciting** for the Deep Learning field, but it’s also **accessible** to many Computer Vision Engineers. I suspect it may even become mandatory inside Deep Learning curriculums in the next few years.

Today, we are going to learn the fundamentals of Multi-Task Learning, and to kick off this article, below is the story of how I first learned about it:

I researched Tesla’s Autopilot a few years ago to write an article. I learned TONS of things through their conferences, articles, and by following the engineers on LinkedIn. But at some point, I stumble across a word:

“**HydraNet**”

HydraNet? Interesting, what is this?

I did some digging and learned that Tesla isn’t using 10 or 20 models. Despite having 20+ tasks to solve, **they’re using just one** (at least, this is the main idea behind that).

They have one model that can solve every possible task they are working on, such as:

*   Object Detection
*   Road curves estimation
*   Depth Estimation
*   3D Reconstruction
*   Video Analysis
*   Object Tracking
*   etc…

![Image 2](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20685%20500'%3E%3C/svg%3E)

**Figure 1:** Tasks needed for a self-driving car ([source](https://www.youtube.com/watch?v=oBklltKXtDE)).

Before we dive into the “how,” I’d like to show you what a HydraNet can do.

Here’s a benchmark of 2 Computer Vision models running on an NVIDIA GPU on 3 different configurations.

*   In the first configuration, we’re running a **semantic segmentation model**.
*   On the second configuration, we’re **_stacking_** a **monocular depth estimation** model.
*   In the third configuration, we’re **building a HydraNet capable of doing both tasks simultaneously**.

Now, look at the results (see **Table 1**):

![Image 3](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20700%20152'%3E%3C/svg%3E)

**Table 1:** Benchmark from Single-Task vs. Multi-Task Learning.

**Semantic segmentation runs at 29 FPS** in the first column and uses **70% of the GPU**. But as soon as **we add Depth Estimation, the GPU becomes full**.

This happened a few years ago while I was working on autonomous shuttles. I trained a YOLOv3 model to detect obstacles and needed to add more tasks, such as lane line estimation.

Suddenly, my GPU was full, and I couldn’t run the car anymore.

So in the third column, you notice that **_the HydraNet saves the GPU usage while keeping the FPS relatively high_**.

Because instead of doing this in parallel:

![Image 4](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20700%20237'%3E%3C/svg%3E)

**Figure 2:** Two Single-Task Models in parallel.

We’re doing this:

![Image 5](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20700%20225'%3E%3C/svg%3E)

**Figure 3:** A Single Multi-Task Model.

This is the same kind of architecture that Tesla uses; as you notice, it’s powerful.

One body, several heads.

![Image 6: Image](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%200%200'%3E%3C/svg%3E)

**Figure 4:** A Hydra.

This is what Tesla is doing, and the article I was referring to is available [here](https://www.thinkautonomous.ai/blog/how-tesla-autopilot-works/).

In this lesson, we will see a simpler version of a multi-task learning project:

* * *

[**Solving a Multi-Task Learning Project**](#TOC)
-------------------------------------------------

Now that you’re familiar with HydraNets, we will have a short tutorial on using HydraNets with PyTorch.

For that, we will use the UTK Face Dataset. It’s a classification dataset where each image has 3 labels:

*   **gender** (male or female)
*   **race** (white, black, Asian, Indian, other)
*   **age** (a continuous number)

**Our HydraNet will therefore have 3 heads, one for each task**. The fun part is that these won’t be the same heads: **predicting the age is a regression task**, but **predicting the gender is a binary classification task** and **predicting the race is a multi-class classification task**.

This all means there will be several loss functions to handle.

![Image 7](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20700%20297'%3E%3C/svg%3E)

**Figure 5:** 3-Task Learning.

With PyTorch, we will create this exact project. For that, we’ll:

1.  Create a **Multi-Task DataLoade**r with PyTorch
2.  Create a **Multi-Task Network**
3.  Train the **Model and Run the Results**

With PyTorch, we always start with a Dataset that we encapsulate in a PyTorch DataLoader and feed to a model.

The DataLoader is our first step.

* * *

### [**Creating a Multi-Task DataLoader with PyTorch**](#TOC)

In this tutorial, I will assume that you are familiar with [PyTorch](https://pyimagesearch.com/?s=pytorch). You know all about datasets and dataloaders, and seeing `__dunders__` doesn’t scare you.

Here is the idea: our data will be fed to a dataset class converted to a PyTorch DataLoader.

![Image 8](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20700%20290'%3E%3C/svg%3E)

**Figure 6:** From Dataset to DataLoader with PyTorch.

On the image:

1.  We collect data, like for any Deep Learning project
2.  We send this data to a custom Dataset class and create a train and test dataset
3.  We encapsulate these two datasets into two PyTorch DataLoaders
4.  When training, we’re working with these DataLoaders

* * *

#### [**Data**](#TOC)

The first is the data.

The dataset is a set of images, and the names of the images give the labels. For example, the image **UTKFace/****100****\_****0****\_****0****\_20170112213500903.jpg.chip.jpg** can be interpreted as follows:

*   **UTKFace/** is a prefix
*   **100** **is the age**
*   **0** **is the gender** (0: male, 1: female)
*   **0** **is the race** (0:White, 1:Black, 2:Asian, 3:Indian, 4:Other)
*   The rest is the date and the extension (jpg)

* * *

#### [**Dataset**](#TOC)

**The Dataset class must implement 3 functions:**

*   `def __init__(self)`: a function defining the elements of a dataset (like inputs and labels)
*   `def __len__(self)`: a function that returns the number of elements
*   `def __getitem__(self)`: a function that returns an item from the dataset

Now, here is how each function is implemented.

* * *

##### [**INIT**](#TOC)

 def \_\_init\_\_(self, image\_paths):
        # Define the Transforms
        self.transform = transforms.Compose(\[transforms.Resize((32, 32)), transforms.ToTensor(), transforms.Normalize(\[0.485, 0.456, 0.406\], \[0.229, 0.224, 0.225\])\])

        # Set Inputs and Labels
        self.image\_paths = image\_paths
        self.images = \[\]
        self.ages = \[\]
        self.genders = \[\]
        self.races = \[\]

        for path in image\_paths:
            filename = path\[8:\].split("\_")
            if len(filename)==4:
                self.images.append(path)
                self.ages.append(int(filename\[0\]))
                self.genders.append(int(filename\[1\]))
                self.races.append(int(filename\[2\]))

In the `init` function, we define the transforms (data augmentation, conversion to a tensor, etc.) operations.

Then, we go over the image paths and take the elements (e.g., `100`, `0`, and `0`), and add them to lists of labels.

Once we have called the `__init__` function, we should have 4 list definitions (images, races, genders, ages).

* * *

##### [**LEN**](#TOC)

This function just returns the number of images.

def \_\_len\_\_(self):
         return len(self.images)

If you didn’t use that, you couldn’t run the training loop for every dataset image. You need the `len` function because `len()` doesn’t work with custom classes; it works with lists, arrays, etc., but not custom elements. So you need to redefine it.

* * *

##### [**GET\_ITEM**](#TOC)

This function will do all the work you defined in `__init__` (e.g., loading an image, calling the transformation, and getting the label). In addition, **this function aims to return a specific element of the dataset given an index.**

def \_\_getitem\_\_(self, index):
        # Load an Image
        img = Image.open(self.images\[index\]).convert('RGB')
        # Transform it
        img = self.transform(img)

        # Get the Labels
        age = self.ages\[index\]
        gender = self.genders\[index\]
        race = self.races\[index\]
        
        # Return the sample of the dataset
        sample = {'image':img, 'age': age, 'gender': gender, 'race':race}
        return sample

**_Note:_** _In this example, we return it as a dictionary, but you could also return 4 elements (this is Python, after all!)_

* * *

#### [**DataLoaders**](#TOC)

Finally, we can have a DataLoader by calling the function:

train\_dataloader = DataLoader(UTKFace(train\_dataset), shuffle=True, batch\_size=BATCH\_SIZE)
val\_dataloader = DataLoader(UTKFace(valid\_dataset), shuffle=False, batch\_size=BATCH\_SIZE)

* * *

### [**Model Definition**](#TOC)

Once we have a DataLoader, **we can define and train our model with PyTorch**.

We could go from scratch, but if you have taken a few Deep Learning courses, you know that **we often prefer to use Transfer Learning**. This is why we’ll start with a pretrained model that we’ll transform to make it match our problem.

![Image 9](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20700%20285'%3E%3C/svg%3E)

**Figure 7:** Backbone for a Multi-Task Model.

The backbone we’ll use can be anything (e.g., a ResNet, MobileNet, Inception, etc.).

**_What matters is the end of the network!_** We want to change the last layer of the pretrained model and, just like transfer learning, adapt it to our problem.

In this case, there are two ways to process:

1.  _Since our case is really simple, we could change the last layer and add 7 neurons:_ **1 for the gender** (binary classification requires one neuron), **5 for the race** (as for 5 classes), **1 for the age** (a regression requires one neuron)
2.  _Or, we could build heads and have several layers for each._

In the first case, we would have something as simple as this:

self.net = models.resnet18(pretrained=True)
self.net.fc = nn.Linear(model.fc.in\_features, 7)

*   On **Line 1**, we download the `resnet18` model from PyTorch.
*   On **Line 2**, notice that we use `model.fc`. The last layer of the pretrained ResNet we download is called “FC” for Fully Connected. We modify that layer by changing the number of neurons. If you had a different model, `fc` could be named something else.

**In the second case, we build heads.** For that, we build sequences of layers on top of the existing FC layer.

We will define:

*   `fc1`: as the fully connected layer for the age
*   `fc2`: for the genders
*   `fc3`: for the races

Here is the code version:

class HydraNet(nn.Module):
    def \_\_init\_\_(self):
        super().\_\_init\_\_()
        self.net = models.resnet18(pretrained=True)
        self.n\_features = self.net.fc.in\_features
        self.net.fc = nn.Identity()
        self.net.fc1 = nn.Sequential(OrderedDict(\[('linear', nn.Linear(self.n\_features,self.n\_features)),('relu1', nn.ReLU()),('final', nn.Linear(self.n\_features, 1))\]))
        self.net.fc2 = nn.Sequential(OrderedDict(\[('linear', nn.Linear(self.n\_features,self.n\_features)),('relu1', nn.ReLU()),('final', nn.Linear(self.n\_features, 1))\]))
        self.net.fc3 = nn.Sequential(OrderedDict(\[('linear', nn.Linear(self.n\_features,self.n\_features)),('relu1', nn.ReLU()),('final', nn.Linear(self.n\_features, 5))\]))
        
    def forward(self, x):
        age\_head = self.net.fc1(self.net(x))
        gender\_head = self.net.fc2(self.net(x))
        race\_head = self.net.fc3(self.net(x))
        return age\_head, gender\_head, race\_head

In the `init()` function, we define sequential models, which are sequences of layers.

_If we consider just the age:_

       self.net.fc1 = nn.Sequential(OrderedDict(\[('linear', nn.Linear(self.n\_features,self.n\_features)),('relu1', nn.ReLU()),('final', nn.Linear(self.n\_features, 1))\]))

We create a dictionary (we don’t have to, but this is good practice to name layers) with a linear Layer of size `(n_features, n_features)`, followed by a ReLU activation and another linear layer of size `(n_features, 1)`.

_Then, in the network:_

       age\_head = self.net.fc1(self.net(x))

Our input goes through the model and then goes through `fc1`.

We then repeat this for the 2 other heads. Again, the forward function returns the 3 heads, but this time, it could also return one dictionary made of the 3 (just like for the dataset class).

Next, we need to define loss functions and train our model.

* * *

#### [**Training: _Loss Functions_**](#TOC)

We’ll begin with the loss functions. In Multi-Task Learning, you still have one loss. The difference is that this loss combines all your other losses.

We could have something as simple as:

![Image 10](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%200%200'%3E%3C/svg%3E)

There are peculiarities to every type of loss. In the gender case, we want to apply a sigmoid activation in the last layer; but in the Race case, we want to apply a softmax.

Now, here is where it will probably fail:

Consider we must predict an age. If our prediction is 50, but the actual age is 30, we have a loss of 20. **If we use Mean Squared Error, our loss becomes (20² = 400).**

![Image 11](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20400%2070'%3E%3C/svg%3E)

Mean squared error loss formula

If we add the classification loss from the genders, it will be a number between 0 and 1.

If we add the two, we’ll end up with a loss affected only by the age error; **our losses aren’t on the same scale**.

A simple way would be to add coefficients in front of the losses and have an overall loss that has balanced values:

![Image 12](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20700%20105'%3E%3C/svg%3E)

An even better way would be to change our age loss from Mean Squared Error to ![Image 13: L_1](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%2016%2015'%3E%3C/svg%3E) Loss:

![Image 14](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20400%2058'%3E%3C/svg%3E)

An error of 20 would stay 20, and we can bring it back to a normal scale.

Finally, the training loop.

* * *

#### [**Training Loop**](#TOC)

We are now ready to train our model with PyTorch!

The code will be minimal to “make it work,” **just for the training**.

net = resnet34(pretrained=True)
model = HydraNet(net).to(device=device)

race\_loss = nn.CrossEntropyLoss() # Includes Softmax
gender\_loss = nn.BCELoss() # Doesn't include Softmax
age\_loss = nn.L1Loss()

optimizer = torch.optim.SGD(model.parameters(), lr=1e-4, momentum=0.09)

Sig = nn.Sigmoid()

All of this is the prelude to our work. We define a model and then:

*   The race loss is the Cross-Entropy Loss
*   The gender loss is the Binary Cross-Entropy Loss
*   The age loss is the ![Image 15: L_1](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%2016%2015'%3E%3C/svg%3E) Loss

Something to note, for binary classification, we’ll use Binary Cross-Entropy that doesn’t include the softmax. Instead, we’ll use **Sigmoid**, more suited for this special case.

_Finally, the loop:_

for epoch in range(n\_epochs):
   model.train()
   total\_training\_loss = 0

    for i, data in enumerate(train\_dataloader):
        inputs = data\["image"\].to(device=device)

        age\_label = data\["age"\].to(device=device)
        gender\_label = data\["gender"\].to(device=device)
        race\_label = data\["race"\].to(device=device)

        optimizer.zero\_grad()

        age\_output, gender\_output, race\_output = model(inputs)
        loss\_1 = race\_loss(race\_output, race\_label)
        loss\_2 = gender\_loss(sig(gender\_output), gender\_label.unsqueeze(1).float())
        loss\_3 = age\_loss(age\_output, age\_label.unsqueeze(1).float())
        loss = loss\_1 + loss\_2 + loss\_3
        loss.backward()
        optimizer.step()
        total\_training\_loss += loss

Then the same goes for validation.

The model will train for a number of epochs, and the loss will be optimized.

Once the model is trained, we can use it for inference and predict our 3 values for any input face!

* * *

### What's next? We recommend [PyImageSearch University](https://pyimagesearch.com/pyimagesearch-university/?utm_source=blogPost&utm_medium=bottomBanner&utm_campaign=What%27s%20next%3F%20I%20recommend).

**Course information:**  
84 total classes • 114+ hours of on-demand code walkthrough videos • Last updated: February 2024  
★★★★★ 4.84 (128 Ratings) • 16,000+ Students Enrolled

**I strongly believe that if you had the right teacher you could _master_ computer vision and deep learning.**

Do you think learning computer vision and deep learning has to be time-consuming, overwhelming, and complicated? Or has to involve complex mathematics and equations? Or requires a degree in computer science?

That’s _not_ the case.

All you need to master computer vision and deep learning is for someone to explain things to you in _simple, intuitive_ terms. _And that’s exactly what I do_. My mission is to change education and how complex Artificial Intelligence topics are taught.

If you're serious about learning computer vision, your next stop should be PyImageSearch University, the most comprehensive computer vision, deep learning, and OpenCV course online today. Here you’ll learn how to _successfully_ and _confidently_ apply computer vision to your work, research, and projects. Join me in computer vision mastery.

**Inside PyImageSearch University you'll find:**

*   ✓ **84 courses** on essential computer vision, deep learning, and OpenCV topics
*   ✓ **84 Certificates** of Completion
*   ✓ **114+ hours** of on-demand video
*   ✓ **Brand new courses released _regularly_**, ensuring you can keep up with state-of-the-art techniques
*   ✓ **Pre-configured Jupyter Notebooks in Google Colab**
*   ✓ Run all code examples in your web browser — works on Windows, macOS, and Linux (no dev environment configuration required!)
*   ✓ Access to **centralized code repos for _all_ 536+ tutorials** on PyImageSearch
*   ✓ **Easy one-click downloads** for code, datasets, pre-trained models, etc.
*   ✓ **Access** on mobile, laptop, desktop, etc.

[Click here to join PyImageSearch University](https://pyimagesearch.com/pyimagesearch-university/?utm_source=blogPost&utm_medium=bottomBanner&utm_campaign=What%27s%20next%3F%20I%20recommend)

* * *

[**Conclusion**](#TOC)
----------------------

Multi-Task Learning is one of the most promising techniques in Deep Learning. Many researchers consider it the future of Artificial Intelligence.

It solves an important speed and memory problem (stacking 20 models can’t be good for your RAM and GPU) and has TONS of benefits when training several tasks. For example, in Computer Vision, researchers noted that co-training specific tasks had better advantages than not training them. Just like learning to play tennis, it makes you very good at Ping-Pong!

If you want to apply these concepts, I recommend [PyImageSearch University](https://pyimagesearch.com/pyimagesearch-university/?utm_medium=BP&utm_campaign=HydraNets-2022-08-17), which contains a HydraNets course for Computer Vision.

In fact, we go much further than this and apply this to semantic segmentation and depth estimation for self-driving cars.

**Here is a preview of what you’ll do:**

![Image 16](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20549%20500'%3E%3C/svg%3E)

**Figure 8:** A HydraNet Project done with Think Autonomous.

To enroll: **[PyImageSearch University](https://pyimagesearch.com/pyimagesearch-university/?utm_medium=BP&utm_campaign=HydraNets-2022-08-17)**

This article was written by Jérémy Cohen from [Think Autonomous](https://www.thinkautonomous.ai/?utm_source=pyimagesearch&utm_medium=blog&utm_campaign=multi_task_learning). Jérémy has been teaching thousands of engineers how to work in cutting-edge fields such as self-driving cars or Advanced Computer Vision through his daily emails and online courses. You can subscribe to Jérémy’s email and read them at [https://www.thinkautonomous.ai/py](https://www.thinkautonomous.ai/py?utm_source=pyimagesearch&utm_medium=blog&utm_campaign=multi_task_learning).

* * *

### [**Citation Information**](#TOC)

**Cohen, J.** “Multi-Task Learning and HydraNets with PyTorch,” _PyImageSearch_, P. Chugh, A. R. Gosthipaty, S. Huot, and R. Raha, eds., 2022, [https://pyimg.co/hdkl7](https://pyimg.co/hdkl7)

@incollection{Cohen\_2022\_HydraNets,
  author = {Jeremy Cohen},
  title = {Multi-Task Learning and HydraNets with PyTorch},
  booktitle = {PyImageSearch},
  editor = {Puneet Chugh and Aritra Roy Gosthipaty and Susan Huot and Ritwik Raha},  year = {2022},
  note = {https://pyimg.co/hdkl7},
}

* * *

![Image 17](data:image/svg+xml,%3Csvg%20xmlns='http://www.w3.org/2000/svg'%20viewBox='0%200%20410%20567'%3E%3C/svg%3E)

#### Join the PyImageSearch Newsletter and Grab My FREE 17-page Resource Guide PDF

Enter your email address below to **join the PyImageSearch Newsletter** and **download my FREE 17-page Resource Guide PDF** on Computer Vision, OpenCV, and Deep Learning.