Week 10: More SegFormer

Hi, welcome back to my blog! My research is basically done, so this week I’m starting something new to explore some more options. I’ve complained previously about using a pre-made SegFormer that doesn’t let me look into how it works. So, this week, I found a SegFormer implementation on PyTorch [1] which I’m going reference and play around with. The goal is to analyze and modify the architecture to specialize it for my project.

Transformer this, transformer that, what even is it?

Now that I have the actual code implementation instead of abstract descriptions, let’s unpack the SegFormer architecture further. I’m using the SegFormer paper [2] as my primary reference, and I’ll try to explain everything as conceptually as I can without going into the equations.

Remember that most image-based neural networks have an Encoder and Decoder. (Most famously, the U-Net) The encoder learns features from the image, and the decoder forms the result.

Encoder: MiT (Mix Transformer)

The MiT was proposed by the original SegFormer paper to be a version of the Vision Transformer tailored for segmentation.

Firstly, overlap patch embeddings converts the image into processable data. Patches of the image are merged into feature maps, similar to a convolution (explained in week 3). The overlapping step is crucial, since it allows the model to preserve spatial continuity

The next steps consist of transformer blocks, which have three sub-steps: efficient self-attention, Mix-FFN, and overlap patch embeddings (again). These steps are done repeatedly in transformer blocks, each step further shrinking the image and creating more features.

In the efficient self-attention step, each pixel is treated as a token, and it essentially looks at the other regions in the image and asks how “important” that region is to understanding what it is. This crucial step allows the transformer to capture the global context throughout the image. Why is it called efficient? Because the SegFormer makes the regions that parttake in self-attention contain multiple pixels, so there are less total queries made between the regions.

Next, the Mix-FFN (Feed-Forward Network) expands and applies a non-linear activation (use a non-linear function, like ReLU, to calculate weights) to the patches, allowing the model to learn local contexts.

Decoder

The decoder uses layers of MLPs (Multi-Layer Perceptrons, they are a type of FFN) to convert the features into an output. First, one MLP is used to unify the channel dimensions of each feature from the encoder (one from each of the four transformer blocks). Then, a second MLP upsamples the features to the same size (H/4 x W/4, where H and W are the height and width of the original image). Afterwards, a third MLP concatenates the four outputs. Finally, the last MLP maps it to the prediction image.

Results

How did our model turn out? Mixed reviews. The model itself worked well in that damage levels were accurately assessed in the rough area of the houses. However, it seems that the model ran into some issues upsampling, so the patches look very blocky. I’ve attempted to upsample and use interpolation, but it the shapes made less sense with that. I believe it’s because the last few upsamples don’t take enough context from the original image. I’m working with my advisor to solve it, but if you have any suggestions or disagree with my analysis, I’d love to hear it!