3D Semantic Segmentation - DIY Self Driving Part 4

Introduction

This is a follow up to Voxel from Multicam and is part of a series where I try to train models to perform common self-driving tasks from scratch.

I’ve previously put together occupancy models for self-driving but that’s only one specific perception task.

Another common driving task is semantic segmentation. Semantic segmentation takes in the image and for every pixel predicts a specific class. This can be used to tell things like walls apart from cars or classify different types of lane lines and curbs on a road.

Classical segmentation operates on images but since we’re operating in 3D it’d be nice to get the same classes as either a birdseye view representation or a voxel representation.

Supervised vs Unsupervised Learning

The occupancy grid models that I’ve made before are trained in an unsupervised manner. This means that there’s generated “ground truth” to train the model on and it instead uses consistency between frames to learn the occupancy grid. This makes it simpler since I don’t need to collect any labels which can be very expensive.

Segmentation is a supervised task. It’s trained by having a ground truth (often human labeled) to compare the model output to. Luckily for us, there’s a number of publicly available datasets that I can pull from.

One dataset that matches our tasks fairly well is BDD100K which has 100k labeled images for things like lane lines and drivable space. For full semantic segmentation it only has about 10k images but should be enough for our purposes.

Auto-labeling Models

How do we take the BDD100K dataset and apply it to a completely different driving dataset to learn a 3D representation?

We can use a pretrained model as an “auto-labeler”. This second model is trained on BDD100K and we can run it on our dataset to generate ground truth labels for our model.

I used a retrained YOLOP for the road surface segmentation and for the semantic segmentation I used a pretrained UPerNet with a ConvNeXt-T backbone.

Using pretrained models on similar tasks made it a lot easier for me to quickly fine-tune the models and start training the 3D model.

Fine Tuning Datasets

BDD100K only has forwards facing dashcam footage as part of its training set so it works well for the forward facing camera in the vehicle but less well for the other cameras such as the backwards facing side repeaters. To augment this dataset we can hand label some data and use it to retrain the model.

In this case I spent a couple of days labeling about 300 images. While 300 images is much less than 100k we can do some tricks to weigh those examples more during training in order to make the model prioritize our examples.

For the labeling software I used Label Studio since it’s open source and fairly easy to get setup with. There’s definitely some rough edges but it has some nifty features. I setup autolabeling within Label Studio so I would only have to fine-tune the results from the model instead of labeling the entire image myself.

Labelling Inference

Since these models are just being used as training data we can cut their computational usage by running them in inference mode and in lower precision.

model = load_semantic_model()
# switch to eval mode
model = model.eval()
# convert the model to fp16 for performance boost and to take up less memory
model = model.half()

# run the model in inference mode to disable autograd tracking
with torch.inference_mode():
    target = model(input)

Voxel Segmentation Model

The previous occupancy model outputted a single probability for each of the points in the voxel grid. For segmentation we need to add in probabilities for each of the various classes. We can do this by adding an extra head to our model decoder. Now we have two final layers, one to predict occupancy probabilities and one to predict the semantic classes for each voxel.

The ground and sky classes have been intentionally omitted so the argmax for those areas is noisy.

To convert this occupancy + semantic classes into an image that we can compare to the autolabeling model we use the same differentiable rendering technique that I used for generating the depth maps. Instead of computing the depths it instead renders out the class probabilities.

Autograd Graph Breaks: 50% Memory Savings!

Rendering all these different camera views and autolabeling models is very memory intensive. Since all of the individual camera losses are independent you can cut down on the memory usage by pre-emptively computing the gradients by calling .backward() on each independent loss.

FlashAttention

FlashAttention is a hip new way of doing multiheaded attention in PyTorch. I decided to port my model over to use BF16 for memory+performance reasons and also switched my model to use flash attention instead of the default PyTorch implementation. This cut my memory usage by ~8% (17522MiB vs 18950MiB) but since the transformer is only a part of the E2E model training it likely had less benefit than on a heavier text transformer model.

TVL1 Loss

I’ve also added a total variation L1 loss from the Neural Volumes paper. This helps reduce per camera artifacts as well as reduces the amount of random noise in the grid outside of the rendered camera view.