Panoptic segmentation on embedded car hardware

Semantic segmentation says “car”. It can’t tell two touching cars apart. Panoptic segmentation adds an id per object, which is what tracking and planning need. Research panoptic models are large, and ours had to share an embedded GPU with everything else in the car. I built it three times between 2019 and 2021.

1. Mask R-CNN, slimmed (TensorFlow, 2019–2020)

The first version in 2019 put Mask R-CNN and a 27-class semantic head on one backbone. Computing instance masks inside tf.data was too slow, so I pre-generated them per image.

In 2020 I rebuilt it for the embedded target: Fast-SCNN as the backbone, Mask R-CNN with ROI-Align, and a unified panoptic head and loss after UPSNet. Input 502×1280, 16 classes, with cars as the object class.

The mask branch had a wiring problem. In training it runs on all region proposals; at inference it should only see the final detections after NMS. The framework couldn’t attach a network part after post-processing, so I wrapped the filtering and NMS inside a plugin that passes everything through in training mode. My own docstring calls it “a clever/stupid hack”. It kept the training graph unchanged and the inference graph small.

Then profiling. I built a tool that freezes the graph, reports milliseconds per inference averaged over 50 runs and writes a Chrome trace. The trace pointed at the ops to cut:

1. 2020-02-25
   Fewer boxes first

   Reordered post-processing so the expensive steps see fewer boxes.

2. 2020-03-30
   Cheaper mask head

   Transposed convolutions became nearest upsampling plus a convolution, and the mask head went from 256 filters to 32. A tile-across-all-objects gather went too.

3. 2020-04-16
   No background masks

   Background detections no longer enter the mask branch.

4. 2020-04-21
   No Python in the pipeline

   A py_func colour-to-class mapping blocked the input pipeline on the GIL; a TensorFlow hash table replaced it.

5. 2020-06-05
   TensorRT 7.1

   align_corners=True everywhere, so TensorRT’s resize matches TensorFlow’s.

At inference it keeps 50 proposals instead of the 100 used in training. Most of the gap between the research model and one that fits the budget was in wiring and op choices, not in the architecture.

2. SOLO from scratch (PyTorch, 2020)

When we moved to PyTorch I switched to SOLO, which predicts a mask per grid cell: no boxes, no anchors, simple post-processing. In August 2020 I implemented SOLO, Decoupled SOLO and SOLOv2 from scratch in one commit of about 2,500 lines: heads, training targets, Matrix NMS, Dice and focal losses, an A2D2 instance reader with 7 classes, and augmentations that keep boxes and masks consistent.

The decoupled head predicts S x-maps and S y-maps and multiplies them, instead of S² full masks. At a 40×40 grid that is 80 output channels instead of 1,600.

The first panoptic model trained in September 2020. Two losses made multi-task training behave:

• Uncertainty weighting with clamps. Learned task weights, with minimum and maximum limits so no task’s weight collapses to zero.
• A consistency loss between the semantic and instance outputs.

The post-processing runs inside the TorchScript graph, so the car’s runtime needs no Python. By December 2020 the panoptic detector ran in our C++ stack on LibTorch.

3. Panoptic FCN and the hybrid (2021)

In April 2021 I reimplemented Panoptic FCN with three heads in 17 days: a faithful one, a custom one with class-agnostic objects, and a class-agnostic one with UNet refinement. The backbone started from my segmentation checkpoint.

Its dynamic kernels don’t export cleanly to TensorRT. Instead of rewriting the model, I split it:

The encoder and FPN go through ONNX and a graph clean-up into an INT8 TensorRT engine: 15 MB on Drive AGX Pegasus. The head, about 2 M parameters, stays TorchScript and runs in LibTorch. Getting there meant removing dynamic linspace and dynamic slicing for TensorRT 8. The hybrid backend went into our C++ engine on 17 October 2021; the engines followed two days later.

Put in TensorRT what TensorRT can build. Leave the rest in TorchScript.

4. One network, three tasks (2022)

The last step merged depth, semantic segmentation and instances into one ResNet-18 network of about 18.6 M parameters, started from my depth-completion checkpoint. The numbers are in the depth completion post.

What I’d keep: profile before slimming, keep post-processing inside the graph, and split a model at the export boundary rather than rewrite it.