FastFCN: Rethinking Dilated Convolution in the Backbone for Semantic Segmentation

In plain words

Segmentation models label every pixel of an image. To keep fine detail, the usual trick is to stop the image network from shrinking the picture, which keeps the computation high. That is slow and uses a lot of memory.

This paper keeps the normal, shrinking backbone and adds a small learned module, Joint Pyramid Upsampling (JPU). JPU rebuilds a sharp, high-resolution feature map from the last three layers of the backbone, using several differently spaced filters to look at context at several scales. The authors report that this reports that accuracy holds while the network runs faster.

TL;DR

Dilated convolutions keep segmentation feature maps sharp but cost compute and memory. FastFCN uses an ordinary backbone and a joint upsampling module, JPU, to rebuild the high-resolution map. Author-reported on PASCAL-Context with ResNet-50, single scale: 51.05 mIoU and 37.56 FPS, against 49.91 mIoU and 18.77 FPS for dilated EncNet tab. 1 tab. 2. The official MIT-licensed code differs from the paper in its losses and learning rates (see Paper vs code).

Problem

A classification backbone downsamples the image 32 times, so the last feature map is too coarse for boundaries §1. Dilated (atrous) convolutions avoid that by removing the last two downsampling steps. Their cost grows with the enlarged maps: with ResNet-101, the 23 residual blocks of the dilated stage need about four times the computation and memory of the original, and 3 further blocks about sixteen times §1 §3.2. The paper targets this cost and assumes that the dilated stage can be replaced without losing accuracy §1.

Method

Reformulation. The paper treats the dilated backbone as a target feature map to approximate. Both the dilated and the stride-based layers can be written with the same regular-convolution stack, applied to differently split inputs §3.1 eq. 2 eq. 3. This makes the job a joint upsampling problem: given a low-resolution map and a high-resolution guide, find a function that maps the guide into the target §3.3 eq. 1. The paper solves it with a CNN instead of iterative optimisation §3.3.3 eq. 4.

Network. The backbone is the standard FCN (ResNet-50 or ResNet-101) with stride convolutions and no dilation §3.2. JPU takes the last three stages, Conv3, Conv4 and Conv5 §3.2. A context module on top of JPU produces the label map: PSP [zhao2017pyramid], ASPP [chen2017rethinking] or Encoding [zhang2018context] §3.2. The paper says JPU's output is at output stride 8 §1.

JPU. Each of the three inputs goes through a regular convolution block that maps it to a common width §3.3.3 fig. 4a. Conv4 and Conv5 are upsampled to the size of Conv3 and concatenated with it §3.3.3 fig. 4b. Four separable convolutions with dilation rates 1, 2, 4 and 8 read this concatenation in parallel §3.3.3. The rate-1 branch captures the relation between the Conv3 map and the rest of the fused map. The rate-2, 4 and 8 branches learn the mapping that produces the stride-reduced Conv4-level output §3.3.3 fig. 5. Their outputs are merged, and a final regular convolution block turns the result into predictions §3.3.3 fig. 4c.

# Joint Pyramid Upsampling (JPU)  [§3.3] [fig. 4] [eq. 4]
# Input:  conv3 (highest resolution), conv4, conv5
# Output: (B, 4d, H3, W3); width d is not stated in the paper

def joint_pyramid_upsampling(conv3, conv4, conv5):
    # Map each input to the same width d with a regular conv block.
    # Code uses a 3x3 kernel, not stated in the paper.
    # [§3.3] [code: encoding/nn/customize.py#L158-L169]
    low3 = conv_block(conv3)              # (B, d, H3, W3)
    low4 = conv_block(conv4)              # (B, d, H4, W4)
    low5 = conv_block(conv5)              # (B, d, H5, W5)
    # Upsample Conv4 and Conv5 to Conv3's size, then concatenate.
    # Code uses bilinear resizing.
    # [§3.3] [code: encoding/nn/customize.py#L185-L190]
    low4 = resize_to(low4, (H3, W3))
    low5 = resize_to(low5, (H3, W3))
    fused = concat([low5, low4, low3])    # (B, 3d, H3, W3)
    # Four separable convs with dilation 1, 2, 4, 8 read the same map.
    # Each branch outputs d channels.
    # [§3.3] [code: encoding/nn/customize.py#L171-L182]
    branch1 = separable_conv(fused, dilation=1)
    branch2 = separable_conv(fused, dilation=2)
    branch4 = separable_conv(fused, dilation=4)
    branch8 = separable_conv(fused, dilation=8)
    # Concatenate the four branches along the channel axis.
    # [code: encoding/nn/customize.py#L190-L192]
    return concat([branch1, branch2, branch4, branch8])

Training loss in the released EncNet scripts. The paper names only pixel-wise cross-entropy §4.1. The released EncNet scripts also add an auxiliary loss and a semantic-encoding loss. They are shown here as a separate step, because the paper does not describe them.

# EncNet training loss as set by the released scripts
# [code: encoding/nn/customize.py#L37-L57] [code: experiments/segmentation/option.py#L38-L43]
# Inputs: main logits, aux logits, SE logits, target labels

def total_loss(main, aux, se_logits, target):
    # Main pixel-wise cross-entropy; the only term the paper names [§4.1]
    loss = cross_entropy(main, target)
    # Auxiliary cross-entropy from the Conv3 branch, weight 0.2 by default
    # [code: encoding/models/encnet.py#L37-L40]
    loss = loss + 0.2 * cross_entropy(aux, target)
    # Binary loss on the classes present in the target, weight 0.2
    se_target = classes_present(target)   # (B, nclass), multi-hot
    loss = loss + 0.2 * bce(sigmoid(se_logits), se_target)
    return loss

Dilated and stride convolutions. The paper's analysis, in one dimension, says a dilated convolution equals a split, a regular convolution and a merge §3.3.2 fig. 3. A stride convolution equals a regular convolution followed by dropping every other element §3.3.2 fig. 3. This analysis gives the reason for the reformulation and is not needed to implement JPU.

Hyperparameters

Values are as stated in the paper unless marked "code". Where the paper says nothing, the entry is "not stated". Code values come from the official repo at code_commit.

Results

Author-reported. Setup: PyTorch §4.1; 4 Titan-Xp GPUs §4.1; dataset and backbone as shown. Baselines are the methods named in each table.

Ablation, PASCAL-Context val, ResNet-50, single scale, 59 classes tab. 1

Speed, Titan-Xp, 512 x 512 input, mean of 100 runs tab. 2

PASCAL-Context val, state of the art, 60 classes with background, multi-scale and flip for ours tab. 3

The abstract reports 53.13% for the same ResNet-101 result [abstract]. The paper's text gives 52.10% (ResNet-50) and 54.03% (ResNet-101) on the 59-class protocol without background §4.3.

ADE20K val tab. 4

ADE20K test, evaluation server tab. 5

The first two entries ranked first and second in the COCO-Place challenge 2017 tab. 5.

Paper vs code

Repository: https://github.com/wuhuikai/FastFCN. The owner is one of the paper's authors (Huikai Wu). The default branch (master) is the pure-Python rewrite at commit 126661b, cited below. The README says the original code that produced the arXiv results is branch v1.0.0 README.md. I did not clone that branch; a WebFetch of its encoding/nn/customize.py (commit 107bd864) shows the same JPU and separable-convolution classes as master, so the JPU discrepancies below apply to both. The rest of v1.0.0 (losses, training script) was not checked. Lines cited here are from master.

• JPU first block. Paper: a regular convolution block, kernel and width not stated §3.3 fig. 4a. Code: 3x3 conv, BN, ReLU, width 512 customize.py base.py.
• Final block. Paper: a further regular convolution block after the separable convolutions turns the features into predictions §3.3 fig. 4c. Code: JPU returns the four-branch concatenation with no further conv customize.py. The next layer is the head's 1x1 conv, 2048 to 512 channels encnet.py.
• Upsampling mode. Paper: not stated. Code: bilinear with align_corners=True base.py.
• Loss. Paper: pixel-wise cross-entropy only §4.1. Code: the EncNet scripts pass --aux --se-loss, adding an auxiliary cross-entropy from the Conv3-level features (weight 0.2 by default, 0.4 in the PSP script) and a binary loss on the classes present (weight 0.2) customize.py encnet.py option.py psp_res50_pcontext.sh.
• Learning rate per module. Paper: a single LR for the network §4.1. Code: 10 times the base LR for the JPU, the head and the auxiliary layer train.py. The backbone uses the base LR.
• Scale augmentation. Paper: scale from 0.5 to 2.0 §4.1. Code: the short side is drawn uniformly from 0.5 x 520 to 2.0 x 520, i.e. 260 to 1040 pixels, then a 480 crop base.py base.py. Code also applies a Gaussian blur with probability 0.5, which the paper does not mention base.py.
• Multi-scale test. Paper: multi-scale and flip for SOTA only; scales not stated tab. 3 footnote. Code: scales 0.5 to 1.75 in steps of 0.25, plus flip, when --ms is set test.py.
• ADE20K fine-tuning. Paper: 120 epochs at LR 0.01, then 20 epochs at 0.001 on train+val §4.3. Code defaults match option.py option.py encnet_res50_ade20k_trainval.sh. The paper does not state the LR schedule for ADE20K; the code uses poly decay train.py.
• FPS. Paper: 512 x 512 input, 100 runs, Titan-Xp tab. 2. Code uses a single-image batch of random data test_fps_params.py.
• Released checkpoints. The README lists EncNet+JPU ResNet-101 at 55.84 (MS) for the ADE20K train+val model, while the paper reports 0.5584 on the test server README.md. These are different splits and should not be compared directly.
• Crop size. The README and paper both say EncNet with ResNet-101 was trained at 576 and EncNet+JPU at 480 to fit 12 GB GPUs §4.3 README.md.

No discrepancies were found for the backbone depth, batch size, crop size, momentum, weight decay, poly schedule power or the PASCAL-Context epoch count.

Searched: the abstract and the paper's links (the project page at wuhuikai.me/FastFCNProject is linked from the README, not the arXiv abstract). The GitHub link in the abstract is the repo above.

Pitfalls

• Missing JPU details. The paper gives neither the width d nor the kernel of the first conv block, nor the BN and ReLU placement. Implement the code's values (width 512, 3x3, BN, ReLU) and note that the paper does not fix them.
• Dilation and padding. In the code each separable branch uses padding equal to its dilation rate, so the spatial size is kept customize.py. Implementing the branches with another padding changes the output size.
• Output resolution. JPU's output has Conv3's resolution, not the input's. The EncNet head then bilinearly upsamples the prediction to the input size encnet.py.
• Auxiliary and SE losses are not in the paper. A faithful reimplementation of the paper's loss is pixel cross-entropy only; the reported numbers come from the released scripts, which include the extra losses customize.py.
• Learning-rate multiplier. The 10 times LR for the JPU and head is a code choice. Without it, the JPU trains from the backbone's LR, which the paper does not discuss train.py.
• Speed measurement. The FPS in the paper compares modules within a common setup. The speedup depends on the head and backbone: FPS for ResNet-101 with JPU versus EncNet is 32.02 against 10.51 tab. 2.
• Computation claims. The "more than three times" reduction in computation is stated in the text §1 §5. The paper gives no FLOP or memory table for it; the FPS table is the only measurement.
• Crop asymmetry. ResNet-101 EncNet was trained at 576 but EncNet+JPU at 480, which the paper attributes to GPU memory §4.3. Results from the two crop sizes are not a clean comparison.
• Class counts. The ablation uses 59 classes without background, and the SOTA table uses 60 with background. The same model gives different mIoU on each protocol tab. 3 footnote §4.3.
• Upsampling in JPU_X. The code also contains a JPU_X variant with a second, stacked module customize.py. It is not in the paper. Do not implement it as part of the described method.

Lineage

• Builds on: dilated convolutions of DeepLab [chen2018deeplab], the PSP pyramid pooling module [zhao2017pyramid], ASPP of DeepLabV3 [chen2017rethinking], and the Encoding module of EncNet [zhang2018context] §2 §3.1.
• Joint upsampling and guided filtering [li2016deep] [wu2018fast] are cited as related, but for 3-channel images; the paper says JPU is built for high-dimensional feature maps §2.
• Compared against: FPN [lin2017feature] tab. 1; bilinear upsampling tab. 1; DUpsampling [tian2019decoders], which the paper says depends on the label space §2 tab. 3.
• Superseded by: not known from content in this wiki.