Perception trilogy: one encoder, three tasks, learned loss weights

Embedded hardware in a car has room for one network, not three. I shared one encoder between depth, segmentation and instances, and let uncertainty weighting, with bounds, balance their losses.

The car needed depth, semantic segmentation and instances. Three networks cost three times the compute on an embedded GPU that also had to run everything else. So from January 2021 I built one network with a shared encoder and a head per task, and called it the “trilogy”: depth, semantic and instance segmentation.

Sharing the encoder was not the hard part. Balancing the losses was.

Every task speaks a different scale

My first multi-task net, in 2018, put segmentation and SSD detection on one encoder. Detection boxes and segmentation masks came from different image sets, so every sample was flagged as one or the other, its loss masked by the flag, and the two losses simply added.

That works for two losses. The trilogy had depth on one side, and on the other a panoptic head with its own mix of mask, position and consistency losses. Each lives on its own scale, and whichever is largest steers the shared encoder. Hand-tuned weights mean a sweep for every config and dataset, and a single loss weight can be very sensitive: in the later depth work, of five velocity-loss weights, only one trained stably.

The real problem is the gradients. All three heads backpropagate into the same encoder, and each task’s share of that update scales with its loss. Depth was the one that suffered: its loss is small next to the semantic and instance losses, so their gradients dominated the encoder and depth hardly moved it.

Without weighting, the two segmentation losses dominate the encoder’s gradient and depth barely moves it. RUW rescales each loss so every task gets a say.

Let the network learn the weights

Kendall et al. (2018) treat each task’s loss weight as a learned noise level σ: a noisy task gets a large σ and a small weight, and a penalty on σ stops the network from switching every task off. Liebel and Körner revised the penalty so it stays positive (RUW, revised uncertainty weighting):

L = Σ (1 / 2σᵢ²) · Lᵢ + ln(1 + σᵢ²)

In September 2020 I added this as a loss module to our PyTorch code base: one learned log-variance per task, optionally normalised. The first user was panoptic segmentation, with four losses: semantic, instance mask, instance class and a consistency term between the two heads.

The panoptic network gave the first clean comparison. Two runs of the same decoupled SOLO model on a ResNet-18, identical except for RUW, both trained for 60 epochs. Its checkpoint metric is [email protected] + mIoU, so higher is better and values above 1 are normal. With RUW it reached 1.357; without it, 1.308.

Same model and schedule; the only change is learned task weights. Best epoch of each 60-epoch run.

I don’t have the same comparison for the three-task network: every trilogy run I kept used RUW.

Then bound them

Learned weights can still drift. If one task’s weight sinks towards zero, the encoder stops learning it. In May 2021 I added a minimum and a maximum weight per task, so no task is ever silenced and none takes over.

Panoptic: minimum weights [0, 0.05, 0.05, 0.2] for its four losses.
Trilogy: minimum [0.4, 0] and maximum [1, 0.6] for the two tasks.

The first version had a bug: it called clamp and threw the result away, so the bounds did nothing. I fixed it the next day, and the loss now returns the clamped weights it actually used, which are the ones I log.

Learn the weights, bound them, and log the ones you actually use.

The trilogy network

One ResNet-18 encoder, three outputs: depth from the depth decoder, semantic classes and instances from the segmentation decoder and its panoptic head. RUW weighs the depth loss against the panoptic loss, within bounds.

The encoder is the RGB-D ResNet-18 from depth completion. The segmentation branch carries a SOLO-style panoptic head with 19 classes and class-agnostic instances.

Datasets that don’t overlap. Depth comes from image sequences with sparse LiDAR; segmentation labels come from single frames. Each training step takes one batch of each. The segmentation batch runs with an empty depth channel, so the encoder learns to work with and without depth.
Fixed weights inside a task, learned weights between tasks. Within the segmentation loss the thing mask, stuff mask and thing position keep hand-set weights (0.045, 0.95, 0.005). RUW only balances depth against segmentation.
How much decoder to share. In February 2021 I tried a fully shared decoder, no shared decoder, and sharing up to the first level. The final config shares the encoder only; each task has its own decoder.
One number to pick a checkpoint: abs_rel + (1 − mIoU) + (1 − [email protected]), lower is better.

Results

The last runs, in January 2022, started from the depth-completion checkpoints and fine-tuned all three tasks together on A2D2.

Starting from the depth checkpoint trained with the moving-object mask brought abs_rel from 0.156 to 0.134, for a small cost in mIoU (0.437 → 0.417) and mAP@50 (0.313 → 0.306).
The best run reached 0.1326 abs_rel, 0.42 mIoU and 0.31 mAP@50 from one ResNet-18. Its depth was slightly better than the depth-only network (0.1354).

To run on the car, the model had to export to TensorRT 8, which did not support a dynamic scale factor in upsampling. The decoders now upsample to an explicit output size, and the whole trilogy exports to TorchScript.

Beyond the car

The same loss went with me into robotics. In 2022 a cube pose network balanced its position and rotation losses with RUW, with the bounds pinning rotation at 0.3 and position between 0.7 and 0.8.