Running a State-of-the-Art Segmentation Model on One GPU

The state of the art needed 8 GPUs

The network I wanted to build on was the Gated Shape CNN (GSCNN) by Takikawa et al., published at ICCV 2019. Its insight is that shape is a different kind of information from color and texture, and that cramming both into one network makes each worse. So GSCNN splits into two streams: a regular stream that encodes color and texture semantics, and a shape stream that encodes edges and boundaries. Gates let the regular stream feed what it knows into the shape stream, and a fusion module recombines them into the final prediction.

The catch was the backbone. GSCNN's regular stream is a WideResNet-38, and the authors trained it on 8 GPUs. Their repository carries a warning that the code will probably not run on a single GPU. I tried anyway, and it didn't. Looking through the issues on the repository, I found other people asking the same thing: was there a version of this that would run on one machine?

There wasn't. As far as I could tell, no publicly accessible implementation swapped in a lighter backbone. That is a real barrier — it means the only people who can experiment with a promising architecture are the ones who already have a cluster.

Rebuilding GSCNN around ResNet-18

So I reimplemented GSCNN with a ResNet-18 regular stream and published it. ResNet-18 has four residual blocks, and those became the structure of the regular stream. The input to its first residual block is also routed into the first residual layer of the shape stream, after a 1x1 convolution that squeezes it down to a single channel.

The shape stream keeps the original three-residual-block structure. Between its blocks — and after the final one — sit the gates, each combining the previous shape-stream layer with a layer of the regular stream (again passed through a 1x1 convolution down to one channel). The first gate takes the shape stream's first block and the regular stream's second; the second gate pairs the shape stream's second block with the regular stream's third; the third gate pairs the shape stream's third block with the regular stream's fourth.

The fusion module is taken directly from the original GSCNN. It runs the combined streams through ASPP and produces both the segmentation loss and the duality task loss, while the shape stream carries its own edge loss.

Making the two streams talk both ways

Here is the part I was actually curious about. In GSCNN, information only flows one way: the shape stream takes layers from the regular stream, because knowing where the objects are helps you find their boundaries. But the reverse ought to be true too — knowing where the semantic edges are should help you decide what the regions between them contain. So why not connect the shape stream back into the regular stream?

To do it, I took the output of the third residual layer of the regular stream through a 1x1 convolution producing 128 channels (128 x h x w), and concatenated it with a matching 128 x h x w tensor from the shape stream along the channel dimension. That yields the 256 x h x w input the fourth residual block already expects, so nothing downstream had to change. That same output is also where I hooked into the shape stream's second gate.

The shape stream gained a duplicate second residual block, so that two identical blocks sit between the first and second gates. The first of those two sends its output up to the regular stream, after a convolutional layer brings it to the 128 x h x w shape needed for the concatenation.

Results

I trained both networks for 20 epochs on the Cityscapes leftImg8bit dataset, which labels 19 classes across urban street scenes, and used it for validation and testing as well. Every hyperparameter and input transform was kept identical to the original GSCNN so that the comparison stayed honest.

The ResNet-18 GSCNN reached its best accuracy of 74.2% at epoch 18. The original GSCNN reports 82.9%, but it got there with 175 epochs on 8 GPUs. My network gives up accuracy, and in exchange it runs on hardware people actually have.

The bidirectional variant reached 48.8%. There is not much to say beyond that: adding the connection made things substantially worse, on both mask quality (mIoU) and boundary quality (F-score). The original topology was clearly well thought out and well tested, and altering it hurt.

What I took away

My hypothesis was wrong, and I think that is worth writing down plainly rather than burying. A negative result is still a result — the intuition that "if A helps B, then B should help A" is appealing and, at least for this connection, it does not hold.

I want to be careful about how far I push that conclusion, though. My experiment was small. I tested one kind of connection, in one place, once. There may well be a version of this idea — a different style of connection, a different number of them, or different locations — that does improve on GSCNN. What I can say is that my experiments give no evidence that this is a meaningful path to go down, and that is useful information for the next person considering it.

The durable output was the other half. There is now a publicly available ResNet-18 version of GSCNN, which is exactly what people on the internet had been asking for, and it trains on a single GPU. My hope is that lowering the hardware bar means more people get to experiment with the architecture — possibly to the point of beating the original.

Thanks are owed to Towaki Takikawa, David Acuna, Varun Jampani, and Sanja Fidler for publishing their GSCNN implementation, and to the staff of MIT 6.819/6.869 for teaching me the computer vision fundamentals this project rests on.