SegFormer and HRNet Comparason for Semantic Segmentation

This (incomplete) repo consists of an image segmentation pipeline on the Cityscapes dataset, using HRNet, and a powerful new transformer-based architecture called SegFormer . The scripts for data preprocessing, training, and inference are done mainly from scratch. The model construction code for HRNet (models/hrnet.py) and SegFormer (models/segformer.py) have been adapted from the official mmseg implementation, whereas models/segformer_simple.py contains a very clean SegFormer implementation that may not be correct.

HRNet and SegFormer are useful architectures to compare, because they represent fundamentally different approaches to image understanding. HRNet - like most other vision architectures - is at its core a series of convolution operations that are stacked, fused, and connected in a very efficient manner. SegFormer, on the other hand, has no convolutional operations, and instead uses transformer layers. It treats each image as a sequence of tokens, where each token represents a 4x4 pixel patch of the image.

For training, the implementation details of the original papers are followed as closely as possible.

Due to memory limitations (single RTX 3090 GPU 24 GB), gradient accumilation was used for training the SegFormer model.

HRNet

SegFormer

Official SegFormer

Replication of the B5 model in the official repository. The number of parameters matches up with the paper. The total number of multiply adds may be irrelevant, since it is difficult to determine if it is the same calculation used in the paper to calculate "flops".

model = Segformer(
    pretrained=cfg.MODEL.PRETRAINED,
    img_size=1024,
    patch_size=4, 
    embed_dims=[64, 128, 320, 512], 
    num_heads=[1, 2, 5, 8], 
    mlp_ratios=[4, 4, 4, 4],
    qkv_bias=True, 
    norm_layer=partial(nn.LayerNorm, eps=1e-6), 
    depths=[3, 6, 40, 3], 
    sr_ratios=[8, 4, 2, 1],
    drop_rate=0.0, 
    drop_path_rate=0.1,
    decoder_dim = 768
)

Total Parameters: 85,915,731

Total Multiply Adds (For Convolution and Linear Layers only): 11,607 GFLOPs

Number of Layers

Conv2d : 107 layers
LayerNorm : 161 layers
OverlapPatchEmbed : 4 layers
Linear : 264 layers
Dropout : 208 layers
Attention : 52 layers
Identity : 2 layers
DWConv : 52 layers
GELU : 52 layers
Mlp : 52 layers
Block : 52 layers
DropPath : 102 layers
LinearMLP : 4 layers
Dropout2d : 1 layers

Simple SegFormer

Code taken from this repo

model = Segformer(
    dims = (64, 128, 320, 512),     # dimensions of each stage
    heads = (1, 2, 5, 8),           # heads of each stage
    ff_expansion = (4, 4, 4, 4),    # feedforward expansion factor of each stage
    reduction_ratio = (8, 4, 2, 1), # reduction ratio of each stage for efficient attention
    num_layers = (3, 6, 40, 3),     # num layers of each stage
    decoder_dim = 768,              # decoder dimension
    num_classes = 19                # number of segmentation classes
).to(device)

Total Parameters: 255,280,531

Total Multiply Adds (For Convolution and Linear Layers only): 679 GFLOPs