Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place. Commercial Alternative to JupyterHub.
Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place. Commercial Alternative to JupyterHub.
Path: blob/main/examples/semantic_segmentation-tf.ipynb
Views: 2535
Fine-tuning for Semantic Segmentation with 🤗 Transformers
This tutorial shows how to fine-tune a SegFormer model in TensorFlow for the task of semantic segmentation. The tutorial is a TensorFlow port of this blog post. As such, the notebook uses code from the blog post.
This notebook shows how to fine-tune a pretrained vision model for Semantic Segmentation on a custom dataset. The idea is to add a randomly initialized segmentation head on top of a pre-trained encoder, and fine-tune the model altogether on a labeled dataset.
Model
This notebook is built for the SegFormer model (TensorFlow variant) and is supposed to run on any semantic segmentation dataset. You can adapt this notebook to other supported semantic segmentation models such as MobileViT.
Data augmentation
This notebook leverages TensorFlow's image module for applying data augmentation. Using other augmentation libraries like albumentations is also supported.
Depending on the model and the GPU you are using, you might need to adjust the batch size to avoid out-of-memory errors. Set those two parameters, then the rest of the notebook should run smoothly.
In this notebook, we'll fine-tune from the https://huggingface.co/nvidia/mit-b0 checkpoint, but note that there are others available on the hub.
Before we start, let's install the datasets, transformers, and evaluate libraries. We also install Git-LFS to upload the model checkpoints to Hub.
|████████████████████████████████| 441 kB 27.3 MB/s
|████████████████████████████████| 5.5 MB 90.7 MB/s
|████████████████████████████████| 72 kB 1.7 MB/s
|████████████████████████████████| 212 kB 88.9 MB/s
|████████████████████████████████| 115 kB 93.0 MB/s
|████████████████████████████████| 163 kB 92.6 MB/s
|████████████████████████████████| 95 kB 6.3 MB/s
|████████████████████████████████| 127 kB 89.3 MB/s
|████████████████████████████████| 7.6 MB 76.1 MB/s
|████████████████████████████████| 115 kB 89.3 MB/s
Error: Failed to call git rev-parse --git-dir --show-toplevel: "fatal: not a git repository (or any of the parent directories): .git\n"
Git LFS initialized.
If you're opening this notebook locally, make sure your environment has an install from the last version of those libraries.
You can share the resulting model with the community. By pushing the model to the Hub, others can discover your model and build on top of it. You also get an automatically generated model card that documents how the model works and a widget that will allow anyone to try out the model directly in the browser. To enable this, you'll need to login to your account.
We also quickly upload some telemetry - this tells us which examples and software versions are getting used so we know where to prioritize our maintenance efforts. We don't collect (or care about) any personally identifiable information, but if you'd prefer not to be counted, feel free to skip this step or delete this cell entirely.
Fine-tuning a model on a semantic segmentation task
In this notebook, we will see how to fine-tune one of the 🤗 Transformers vision models on a Semantic Segmentation dataset.
Given an image, the goal is to associate each and every pixel to a particular category (such as table). The screenshot below is taken from a SegFormer fine-tuned on ADE20k - try out the inference widget!
Loading the dataset
We will use the 🤗 Datasets library to download our custom dataset into a DatasetDict.
We're using the Sidewalk dataset which is dataset of sidewalk images gathered in Belgium in the summer of 2021. You can learn more about the dataset here.
Let us also load the Mean IoU metric, which we'll use to evaluate our model both during and after training.
IoU (short for Intersection over Union) tells us the amount of overlap between two sets. In our case, these sets will be the ground-truth segmentation map and the predicted segmentation map. To learn more, you can check out this article.
The ds object itself is a DatasetDict, which contains one key per split (in this case, only "train" for a training split).
Here, the features tell us what each example is consisted of:
pixel_values: the actual imagelabel: segmentation mask
To access an actual element, you need to select a split first, then give an index:
Each of the pixels above can be associated to a particular category. Let's load all the categories that are associated with the dataset. Let's also create an id2label dictionary to decode them back to strings and see what they are. The inverse label2id will be useful too, when we load the model later.
Note: This dataset specificaly sets the 0th index as being unlabeled. We want to take this information into consideration while computing the loss. Specifically, mask the pixels for which the network predicted unlabeled and don't compute loss for it since they don't contribute to training that much.
Preprocessing the data
Before we can feed these images to our model, we need to preprocess them.
Preprocessing images typically comes down to (1) resizing them to a particular size (2) normalizing the color channels (R,G,B) using a mean and standard deviation. These are referred to as image transformations.
To make sure we (1) resize to the appropriate size (2) use the appropriate image mean and standard deviation for the model architecture we are going to use, we instantiate what is called an image processor with the AutoImageProcessor.from_pretrained method.
This image processor is a minimal preprocessor that can be used to prepare images for model training and inference.
Next, we can preprocess our dataset by applying these functions. We will use the set_transform functionality, which allows to apply the functions above on-the-fly (meaning that they will only be applied when the images are loaded in RAM).
Training the model
Now that our data is ready, we can download the pretrained model and fine-tune it. We will the TFSegformerForSemanticSegmentation class. Calling the from_pretrained method on it will download and cache the weights for us. As the label ids and the number of labels are dataset dependent, we pass label2id, and id2label alongside the model_checkpoint here. This will make sure a custom segmentation head will be created (with a custom number of output neurons).
The warning is telling us we are throwing away some weights (the weights and bias of the decode_head layer) and randomly initializing some other (the weights and bias of a new decode_head layer). This is expected in this case, because we are adding a new head for which we don't have pretrained weights, so the library warns us we should fine-tune this model before using it for inference, which is exactly what we are going to do.
Note that most models on the Hub compute loss internally, so we actually don't have to specify anything there! Leaving the loss field blank will cause the model to read the loss head as its loss value.
This is an unusual quirk of TensorFlow models in 🤗 Transformers, so it's worth elaborating on in a little more detail. All 🤗 Transformers models are capable of computing an appropriate loss for their task internally (for example, a CausalLM model will use a cross-entropy loss). To do this, the labels must be provided in the input dict (or equivalently, in the columns argument to to_tf_dataset()), so that they are visible to the model during the forward pass.
This is quite different from the standard Keras way of handling losses, where labels are passed separately and not visible to the main body of the model, and loss is handled by a function that the user passes to compile(), which uses the model outputs and the label to compute a loss value.
The approach we take is that if the user does not pass a loss to compile(), the model will assume you want the internal loss. If you are doing this, you should make sure that the labels column(s) are included in the input dict or in the columns argument to to_tf_dataset.
If you want to use your own loss, that is of course possible too! If you do this, you should make sure your labels column(s) are passed like normal labels, either as the second argument to model.fit(), or in the label_cols argument to to_tf_dataset.
We need to convert our datasets to a format Keras understands. The easiest way to do this is with the to_tf_dataset() method. Note that our data collators are designed to work for multiple frameworks, so ensure you set the return_tensors='tf' argument to get TensorFlow tensors out - you don't want to accidentally get a load of torch.Tensor objects in the middle of your nice TF code!
train_set is now a tf.data.Dataset type object. We see that it contains two elements - labels and pixel_values. :
The last thing to define is how to compute the metrics from the predictions. We need to define a function for this, which will just use the metric we loaded earlier. The only preprocessing we have to do is to take the argmax of our predicted logits.
In addition, let's wrap this metric computation function in a KerasMetricCallback. This callback will compute the metric on the validation set each epoch, including printing it and logging it for other callbacks like TensorBoard and EarlyStopping.
Why do it this way, though, and not just use a straightforward Keras Metric object? This is a good question - on this task, metrics such as Accuracy are very straightforward, and it would probably make more sense to just use a Keras metric for those instead. However, we want to demonstrate the use of KerasMetricCallback here, because it can handle any arbitrary Python function for the metric computation.
How do we actually use KerasMetricCallback? We simply define a function that computes metrics given a tuple of numpy arrays of predictions and labels, then we pass that, along with the validation set to compute metrics on, to the callback:
Now we can train our model. We can also add a callback to sync up our model with the Hub - this allows us to resume training from other machines and even test the model's inference quality midway through training! Make sure to change the username if you do. If you don't want to do this, simply remove the callbacks argument in the call to fit().
Alternatively, we can also leverage metric to compute different quantities on the validation set in a batchwise manner: mean_iou, mean_accuracy, etc.
Inference
Now that the fine-tuning is done, we can perform inference with the model. In this section, we will also compare the model predictions with the ground-truth labels. This comparison will help us determine the plausible next steps.
Our model is not perfect but it's getting there. Here are some ways to make it better:
Our dataset is small and adding more samples to it will likely help the model.
We didn't perform any hyperparameter tuning for the model. So, searching for better hyperparameters could be helpful.
Finally, using a larger model for fine-tuning could be beneficial.
You can now share this model with all your friends, family, favorite pets: they can all load it with the identifier "your-username/the-name-you-picked" so for instance: