Transfer Learning for Segmentation Using DeepLabv3 in PyTorch | by Manpreet Singh Minhas | Towards Data Science

Back when I was researching segmentation using Deep Learning and wanted to run some experiments on DeepLabv3[1] using PyTorch, I couldn’t find any online tutorial. What added to the challenge was that torchvision not only does not provide a Segmentation dataset but also there is no detailed explanation available for the internal structure of the DeepLabv3 class. However, I did the transfer learning on my own, and want to share the procedure so that it may potentially be helpful for you.
In this article, I’ll be covering how to use a pre-trained semantic segmentation DeepLabv3 model for the task of road crack detection in PyTorch by using transfer learning. The same procedure can be applied to fine-tune the network for your custom dataset. If you want to look at the results and repository link directly, please scroll to the bottom.
Let us start with a brief introduction to image segmentation. The primary goal of a segmentation task is to output pixel-level output masks in which regions belonging to certain categories are assigned the same distinct pixel value. If you color-code these segmentation masks by assigning a different color for every category for visualizing them, then you’ll get something like an image from a coloring book for kids. An example is shown below.
Segmentation has existed for a very long time in the domain of Computer Vision and Image processing. Some of the techniques are simple thresholding, clustering based methods such as k means clustering-segmentation, region growing methods, etc. [3]
With recent advancements in deep learning and the success of convolutional neural networks in image-related tasks over the traditional methods, these techniques have also been applied to the task of image segmentation.
One of these network architectures is DeepLabv3 by Google. Explaining how the model works is beyond the scope of this article. Instead, we shall focus on how to use a pre-trained DeepLabv3 network for our data-sets. We will discuss transfer learning briefly for this.
Deep Learning models tend to struggle when limited data is available. And for most practical applications, getting access to a copious dataset can be very difficult if not impossible. Annotation is tedious and time-consuming. Even if you plan to outsource it, you would still have to shell out money. Efforts have been made to be able to train models from limited data. And of such techniques is called Transfer Learning.
Transfer learning involves the use of a network pre-trained for a source domain and task (in which you hopefully have access to a large dataset) and adopting it for your intended/target domain and task (that is similar to the original task and domain)[4]. It can be conceptually represented by the following diagram.
We change the target segmentation sub-network as per our own requirements and then either train a part of the network or the entire network.The learning rate chosen is lower than in the case of normal training. This is because the network already has good weights for the source task. We don’t want to change the weights too much too fast.Also sometimes the initial layers can be kept frozen since it is argued that these layers extract general features and can be potentially used without any changes.
We change the target segmentation sub-network as per our own requirements and then either train a part of the network or the entire network.
The learning rate chosen is lower than in the case of normal training. This is because the network already has good weights for the source task. We don’t want to change the weights too much too fast.
Also sometimes the initial layers can be kept frozen since it is argued that these layers extract general features and can be potentially used without any changes.
If you want to learn more about Transfer Learning, I talk about it in detail in my paper: Anomaly Detection in Images [4]. Next, I’ll talk about the dataset being used in this article before proceeding with the PyTorch related sections.
For this tutorial, I’ll be using the CrackForest[5][6] dataset for road crack detection using segmentation. It consists of urban road surface images with cracks as defects. The images contain confounding regions such as shadows, oil spills, and water stains. The images were taken using an ordinary iPhone5 camera. The dataset contains 118 images and has corresponding pixel level masks for the cracks, all having a size of 320×480. The additional confounders along with the limited number of samples available for training make CrackForest a challenging dataset [7].
Let us begin by constructing a dataset class for our model which will be used to get training samples. For segmentation, instead of a single valued numeric label that could be one hot encoded, we have a ground truth mask image as the label. The mask has pixel level annotations available as shown in Fig. 3. Therefore, the training tensors for both input and labels would be four dimensional. For PyTorch, these would be: batch_size x channels x height x width.
We will be defining our segmentation dataset class now. The class definition is as follows.
We use the VisionDataset class from torchvision as the base class for the Segmentation dataset. Following three method need to be overloaded.
__init__: This method is where the dataset object would be initialized. Usually, you need to build your image file paths and corresponding labels which are mask file paths for segmentation. These paths are then used in the __len__ and __getitem__ method.__getitem__: This method is called whenever you would use object[index] to access any element. So we need to write the image and mask loading logic here. So in essence, you get one training sample from your dataset using dataset object from this method.__len__: This method is invoked whenever len(obj) would be used. This method simply returns the number of the training samples in the directory.
__init__: This method is where the dataset object would be initialized. Usually, you need to build your image file paths and corresponding labels which are mask file paths for segmentation. These paths are then used in the __len__ and __getitem__ method.
__getitem__: This method is called whenever you would use object[index] to access any element. So we need to write the image and mask loading logic here. So in essence, you get one training sample from your dataset using dataset object from this method.
__len__: This method is invoked whenever len(obj) would be used. This method simply returns the number of the training samples in the directory.
When creating your custom datasets for PyTorch please remember to use PIL library. This allows you to use the torchvision transforms directly without having to define your own.In the first version of this class, I was using OpenCV for just loading the image! (I would’ve met past me, I would have advised myself not to do so.) The library is not only extremely heavy but also not compatible with the torchvision transforms. I had to write my custom transforms and handle dimensional changes myself.
When creating your custom datasets for PyTorch please remember to use PIL library. This allows you to use the torchvision transforms directly without having to define your own.
In the first version of this class, I was using OpenCV for just loading the image! (I would’ve met past me, I would have advised myself not to do so.) The library is not only extremely heavy but also not compatible with the torchvision transforms. I had to write my custom transforms and handle dimensional changes myself.
I’ve added additional functionality to allow you to keep your dataset in a single directory instead of splitting the Train, Val into separate folders because a lot of datasets I was using weren’t in this format and I didn’t want to restructure my folder structure every time.
Now that we have the dataset class defined, the next step is to create a PyTorch data loader from this. Data loaders allow you to create batches of data samples and labels using multiprocessing. This makes the data loading process much faster and efficient. The DataLoader class available under torch.utils.data is used for this purpose. The creation process itself is straightforward. You create a DataLoader object by passing the dataset object to it. The supported arguments are shown below.
DataLoader(dataset, batch_size=1, shuffle=False, sampler=None,           batch_sampler=None, num_workers=0, collate_fn=None,           pin_memory=False, drop_last=False, timeout=0,           worker_init_fn=None, *, prefetch_factor=2,           persistent_workers=False)
Few useful arguments are explained below:
dataset (Dataset): dataset from which to load the data. batch_size (int, optional): how many samples per batch to load (default: 1)
batch_size (int, optional): how many samples per batch to load (default: 1).
shuffle (bool, optional): set to True to have the data reshuffled at every epoch. (default: False)
num_workers (int, optional): how many subprocesses to use for data loading. 0 means that the data will be loaded in the main process. (default: 0) Tip: You can set this value equal to the number of cores in your system’s processor as the optimal value. Setting a higher value may lead to degradation in performance.
Additionally, I’ve written two helper functions that give you data loaders depending on your data directory structure and are available in datahandler.py file.
get_dataloader_sep_folder: Create Train and Test data loaders from two separate Train and Test folders. The directory structure should be as follows.
get_dataloader_sep_folder: Create Train and Test data loaders from two separate Train and Test folders. The directory structure should be as follows.
data_dir --Train ------Image ---------Image1 ---------ImageN ------Mask ---------Mask1 ---------MaskN --Train ------Image ---------Image1 ---------ImageN ------Mask ---------Mask1 ---------MaskN
2. get_dataloader_single_folder: Create from a single folder. The structure should be as follows.
--data_dir------Image---------Image1---------ImageN------Mask---------Mask1---------MaskN
These give you training and validation data loaders that shall be used in the training process.
Next, we discuss the crux of this tutorial i.e. how to load a pre-trained model and change the segmentation head according to our data requirements.
Torchvision has pre-trained models available and we shall be using one of those models. I’ve written the following function which gives you a model that has a custom number of output channels. You can change this value if you have more than one class.
First, we get the pre-trained model using the models.segmentation.deeplabv3_resnet101 method that downloads the pre-trained model into our system cache. Note resnet101 is the backbone for the deeplabv3 model obtained from this particular method. This decides the feature vector length that is passed onto the classifier.The second step is the major step of modifying the segmentation head i.e. the classifier. This classifier is the part of the network and is responsible for creating the final segmentation output. The change is done by replacing the classifier module of the model with a new DeepLabHead with a new number of output channels. 2048 is the feature vector size from the resnet101 backbone. If you decide to use another backbone, please change this value accordingly.Finally, we set the model is set to train mode. This step is optional since you can also do this in the training logic.
First, we get the pre-trained model using the models.segmentation.deeplabv3_resnet101 method that downloads the pre-trained model into our system cache. Note resnet101 is the backbone for the deeplabv3 model obtained from this particular method. This decides the feature vector length that is passed onto the classifier.
The second step is the major step of modifying the segmentation head i.e. the classifier. This classifier is the part of the network and is responsible for creating the final segmentation output. The change is done by replacing the classifier module of the model with a new DeepLabHead with a new number of output channels. 2048 is the feature vector size from the resnet101 backbone. If you decide to use another backbone, please change this value accordingly.
Finally, we set the model is set to train mode. This step is optional since you can also do this in the training logic.
So far we’ve covered how to create the data loaders and the DeepLabv3 model with the modified head.
The next step is to train the model.
I’ve defined the following train_model function that trains the model. It saves the training and validation loss and metric (if specified) values into a CSV log file for easy access. The training code is as follows. It is well documented to explain what is going on.
However, I’ll share a few things to keep in mind.
Ensure that you send the model as well as your inputs and label to the same device (which would be cpu or cuda).Remember to clear the gradients using optimizer.zero_grad() before you do forward and backward propagation and parameter update.When training, set the model to train mode by using mode.train()When doing inference, set the model to eval mode by using mode.eval(). This is really important since this ensures that the network parameters are adjusted to account for techniques such as batch norm, dropouts etc. which affect the network weights.
Ensure that you send the model as well as your inputs and label to the same device (which would be cpu or cuda).
Remember to clear the gradients using optimizer.zero_grad() before you do forward and backward propagation and parameter update.
When training, set the model to train mode by using mode.train()
When doing inference, set the model to eval mode by using mode.eval(). This is really important since this ensures that the network parameters are adjusted to account for techniques such as batch norm, dropouts etc. which affect the network weights.
The best model is decided by the lowest loss value. You can select the best model based on an evaluation metric too. But you’ll have to modify the code a bit.
I’ve used the mean squared error (MSE) loss function for this task. The reason I used MSE is that it is a simple function, provides better results and gives a better surface for calculating the gradients. The loss is calculated at the pixel level in our case and is defined as follows:
Note: This loss function won’t work for datasets that have more than one channel. For those cases, you will have to use cross-entropy instead. https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html
I’ve used the Adadelta optimizer with the default settings and a learning rate of 0.0001. The training was done for 25 epochs.
To evaluate the quantitative performance of the models, two metrics were selected. The first metric was the area undercurve (AUC) measurement of the receiver operating characteristics (ROC) [8]. AUC or AUROC is a reliable measure of the degree or measure of the separability of any binary classifier (binary segmentation masks in this case). It provides an aggregate measure of the model’s performance across all possible classification thresholds. An excellent model has AUROC value near to the one and it means that the classifier is virtually agnostic to the choice of a particular threshold. The second metric used for the assessment was the F1 score. It is defined as the harmonic mean of precision (P) and recall (R) and is given by the following equation.
F1 score reaches its best value at one and the worst score at zero. It is a robust choice for classification tasks since it takes both the false positives and false negatives into account.
The best model achieved a testing AUROC value of 0.842. This is a great score and is also reflected by segmentation output obtained after a thresholding operation.
I’ve shown the loss and evaluation metrics during the training in the following graph.
We can observe that the loss value decreased throughout the training. The AUROC and F1 Score values improved with the training. However, we see that the F1 Score values are consistently lower for both training and validation. In fact, these are poor values. The reason for this observation is that I used a threshold of 0.1 for calculating this metric. This wasn’t selected based on the dataset. F1 Score value can change depending on the threshold. However, AUROC is a robust metric that takes into account all possible thresholds. So when you have a binary classification task, it is advisable to use the AUROC metric. Even though the model is performing well on the dataset, as can be seen from the Segmentation Output image, the masks are over dilated in comparison to the ground truth. Maybe since the model is deeper than required, we are observing this kind of behavior. If you have any comments about this phenomenon, please do comment, I would love to know your thoughts.
With this, we have reached the end of the tutorial!
We learnt how to do transfer learning for the task of semantic segmentation using DeepLabv3 in PyTorch on our custom dataset.
First we gained understanding about image segmentation and transfer learning.Next, we saw how to create the dataset class for segmentation for training the model.This was followed by the most important step of how to change the segmentation head of the DeepLabv3 model as per our dataset.The approach was tested for road crack detection on the CrackForest dataset. It achieved an impressive AUROC score of 0.842 after just 25 epochs.
First we gained understanding about image segmentation and transfer learning.
Next, we saw how to create the dataset class for segmentation for training the model.
This was followed by the most important step of how to change the segmentation head of the DeepLabv3 model as per our dataset.
The approach was tested for road crack detection on the CrackForest dataset. It achieved an impressive AUROC score of 0.842 after just 25 epochs.
The code is available at https://github.com/msminhas93/DeepLabv3FineTuning.
Thank you for reading the article. Hope you learned something new from this article. If you found this useful please consider citing my work.
Bibtex Entry:
@misc{minhas_2019, title={Transfer Learning for Semantic Segmentation using PyTorch DeepLab v3}, url={https://github.com/msminhas93/DeepLabv3FineTuning}, journal={GitHub.com/msminhas93}, author={Minhas, Manpreet Singh}, year={2019}, month={Sep}}
IEEE Format Citation:
M. S. Minhas, “Transfer Learning for Semantic Segmentation using PyTorch DeepLab v3,” GitHub.com/msminhas93, 12-Sep-2019. [Online]. Available: https://github.com/msminhas93/DeepLabv3FineTuning.
[1] Rethinking Atrous Convolution for Semantic Image Segmentation, arXiv:1706.05587, Available: https://arxiv.org/abs/1706.05587
[2] Encoder-Decoder with Atrous Separable Convolution for Semantic Image Segmentation, arXiv:1802.02611, Available: https://arxiv.org/abs/1802.02611
[3] https://scikit-image.org/docs/dev/user_guide/tutorial_segmentation.html
[4] Anomaly Detection in Images, arXiv:1905.13147, Available: https://arxiv.org/abs/1905.13147
[5] Yong Shi, Limeng Cui, Zhiquan Qi, Fan Meng, and Zhensong Chen. Automatic road crack detection using randomstructured forests.IEEE Transactions on Intelligent Transportation Systems, 17(12):3434–3445, 2016.
[6] https://github.com/cuilimeng/CrackForest-dataset
[7] AnoNet: Weakly Supervised Anomaly Detection in Textured Surfaces, arXiv:1911.10608, Available: https://arxiv.org/abs/1911.10608
[8] Charles X. Ling, Jin Huang, and Harry Zhang. Auc: A statistically consistent and more discriminating measurethan accuracy. InProceedings of the 18th International Joint Conference on Artificial Intelligence, IJCAI’03,pages 519–524, San Francisco, CA, USA, 2003. Morgan Kaufmann Publishers Inc.