HarDNet-MSEG: A Simple Encoder-Decoder Polyp Segmentation Neural Network that Achieves over 0.9 Mean Dice and 86 FPS

In plain words

Polyps are small growths in the colon that can turn into cancer. A doctor looking at a colonoscopy video has to spot them by eye, and older image methods miss many of them. Polyp segmentation means labelling, pixel by pixel, which part of an image is polyp and which is not.

HarDNet-MSEG is a neural network with two parts. The first part, the encoder, reads the image with a fast, memory-friendly network called HarDNet68. The second part, the decoder, takes features from three levels of that encoder, each at a different zoom, and merges them into one label map. The parts are chosen for speed, and the design avoids attention modules.

The authors report that it is more accurate and faster than the networks they compare with on five public datasets. On Kvasir-SEG it reaches a mean Dice of 0.904 at 86.7 frames per second on one GeForce RTX 2080 Ti GPU. Dice measures how much the predicted mask overlaps the true mask, where 1 means a perfect overlap.

TL;DR

HarDNet-MSEG is a polyp segmentation network with a HarDNet68 encoder and a decoder built from RFB blocks and a dense aggregation step. The decoder drops the highest-resolution features and uses no attention modules §3, §5. Under the Kvasir-SEG 880/120 split it reports 0.904 mean Dice at 86.7 FPS, and under the PraNet protocol it reports 0.912 mean Dice tab. 1 tab. 2 §5. Those figures are author-reported and use two different training setups. Several training details that matter for reproduction (loss, optimizer for each setup, checkpoint selection) are either missing from the paper or differ from the code.

Problem

Manual polyp detection depends on the endoscopist's skill, and earlier feature-based segmentation methods miss many polyps §1. Deep U-Net-style networks and their variants do well on most polyps, but they still cut off boundaries, miss small polyps and break up large regions. They are also slow to run and slow to train §1. The authors aim for a network that is accurate on polyp benchmarks and fast at inference, without extra attention or context modules §1 §5.

Method

HarDNet-MSEG has an encoder (HarDNet68 backbone) and a decoder §3. The decoder follows the design of the Cascaded Partial Decoder §3.2. The paper describes the components in prose and figures; the pseudocode below adds the concrete layer choices read from the released code, and each such line cites the code.

Backbone §3.1. HarDNet is a dense-style CNN whose blocks use fewer shortcut connections than DenseNet, wider key layers, and a few 1x1 convolutions. The authors say this lowers memory traffic and inference time compared with DenseNet and ResNet §3.1. They use HarDNet68 as the backbone because FC-HarDNet70 is strong on Cityscapes §3.1. The decoder does not use HarDBlocks §3.2.

Decoder §3.2. The paper argues that shallow features are high-resolution and expensive, and that deeper features already carry spatial detail, so the decoder discards the shallow features and computes on the deeper ones §3.2. Each kept level passes through an RFB block §3.2.1 fig. 4. The levels are then fused by upsampling and element-wise multiplication §3.2.2 fig. 5.

Metrics §4.3. The paper gives the following formulas. They are unnumbered in the source.

\mathrm{Dice} = \frac{2\,tp}{2\,tp + fp + fn} §4.3

\mathrm{mIoU} = \frac{tp}{tp + fp + fn} §4.3

\mathrm{Recall} = \frac{tp}{tp + fn}, \qquad \mathrm{Precision} = \frac{tp}{tp + fp} §4.3

F_2 = \frac{5\,p\,r}{4p + r} §4.3

\mathrm{Acc.} = \frac{tp + tn}{tp + tn + fp + fn} §4.3

Model forward pass, as implemented in the released code:

# HarDNet-MSEG forward pass  [§3]  [code: lib/HarDMSEG.py#L142-L160]
# Input:  images (B, 3, H, W)
# Output: logits (B, 1, H, W)

def forward(images):
    # Encode with HarDNet68 and keep four feature maps  [§3.1]
    # Strides 4, 8, 16, 32; channels 128, 320, 640, 1024
    # [code: lib/hardnet_68.py#L186-L189]
    _, feat8, feat16, feat32 = hardnet68(images)
    # Reduce each kept map to 32 channels with an RFB block  [§3.2.1] [fig. 4]
    # [code: lib/HarDMSEG.py#L152-L154]
    small8 = rfb8(feat8)          # (B, 32, H/8, W/8)
    small16 = rfb16(feat16)       # (B, 32, H/16, W/16)
    small32 = rfb32(feat32)       # (B, 32, H/32, W/32)
    # Fuse the three levels, deepest first  [§3.2.2] [fig. 5]
    # [code: lib/HarDMSEG.py#L156]
    coarse = dense_aggregation(small32, small16, small8)   # (B, 1, H/8, W/8)
    # Upsample by 8 to the input resolution  [code: lib/HarDMSEG.py#L158]
    return upsample_bilinear(coarse, scale=8)

RFB block, as implemented in the released code (the "RFB" of §3.2.1):

# RFB block  [§3.2.1] [fig. 4]  [code: lib/HarDMSEG.py#L27-L63]
# Input: x (B, C_in, h, w); output (B, 32, h, w)
# conv_bn = conv + batch norm, no ReLU  [code: lib/HarDMSEG.py#L15-L24]

def rfb(x):
    # Branch 0: a single 1x1 conv  [code: lib/HarDMSEG.py#L31-L33]
    branches = [conv_bn(x, kernel=1)]
    # Branches 1 to 3: 1x1 conv, then 1xk and kx1 convs, then a 3x3 conv
    # with dilation k, for k = 3, 5, 7  [code: lib/HarDMSEG.py#L34-L51]
    for k in [3, 5, 7]:
        y = conv_bn(x, kernel=1)
        y = conv_bn(y, kernel=(1, k))
        y = conv_bn(y, kernel=(k, 1))
        branches.append(conv_bn(y, kernel=3, dilation=k))
    # Concatenate the four branches and fuse with a 3x3 conv  [code: L52, L60]
    fused = conv_bn(concat(branches), kernel=3)
    # Add a 1x1 projection of the input, then ReLU  [code: L53, L62]
    return relu(fused + conv_bn(x, kernel=1))

Dense aggregation, as implemented in the released code:

# Dense aggregation of three 32-channel levels  [§3.2.2] [fig. 5]
# [code: lib/HarDMSEG.py#L85-L100]
# deep is the coarsest map, shallow the finest

def dense_aggregation(deep, middle, shallow):
    # Upsample the deep map and multiply it into the middle level  [code: L87]
    middle_mix = conv_a(upsample(deep)) * middle
    # Upsample the deep map twice and the middle map once; multiply all
    # three into the shallow level  [code: lib/HarDMSEG.py#L88-L89]
    shallow_mix = (conv_b(upsample(upsample(deep)))
                   * conv_c(upsample(middle))
                   * shallow)
    # Concatenate with the upsampled deep map, then fuse  [code: L91-L92]
    middle_out = conv_concat2(
        concat(middle_mix, conv_d(upsample(deep))))
    # Upsample the middle output to the shallow level, fuse  [code: L94-L95]
    shallow_out = conv_concat3(
        concat(shallow_mix, conv_e(upsample(middle_out))))
    # A 3x3 conv, then a 1x1 conv down to one channel  [code: L97-L98]
    return conv_1x1(conv_3x3(shallow_out))

Training step, as implemented in the released code. The paper does not state the loss, the batch size or the multi-scale schedule; see the Hyperparameters and Paper vs code sections.

# One training batch  [code: Train.py#L73-L99]
# images (B, 3, T, T) and masks (B, 1, T, T), with T = trainsize

def train_batch(images, masks):
    # Train at three scales: 0.75x, 1x and 1.25x of T, each rounded to a
    # multiple of 32  [code: Train.py#L73, L83]
    for scale in [0.75, 1.0, 1.25]:
        size = round(T * scale / 32) * 32
        scaled_images = resize_bilinear(images, size)
        scaled_masks = resize_bilinear(masks, size)
        # Clear gradients, run the model and compute the loss  [code: L77-L91]
        optimizer.zero_grad()
        logits = model(scaled_images)
        loss = structure_loss(logits, scaled_masks)
        # Backprop, then clamp each gradient element to +/-0.5  [code: L97-L98]
        loss.backward()
        clamp_gradients(optimizer, 0.5)
        # One optimizer step per scale  [code: Train.py#L99]
        optimizer.step()

Loss used by the code (not described in the paper):

# Structure loss: weighted BCE plus weighted soft IoU  [code: Train.py#L14-L25]
# Computed per image over height and width, then averaged over the batch

def structure_loss(logits, masks):
    # Pixels near mask edges get a larger weight: 1 + 5 * |local mean - mask|
    # with a 31x31 window  [code: Train.py#L16]
    weight = 1 + 5 * abs(avg_pool(masks, 31, stride=1, pad=15) - masks)
    # Weighted binary cross-entropy, normalised by the total weight  [L17-L18]
    wbce = sum(weight * bce_with_logits(logits, masks)) / sum(weight)
    probs = sigmoid(logits)
    # Weighted soft IoU  [code: Train.py#L20-L23]
    inter = sum(probs * masks * weight)
    union = sum((probs + masks) * weight)
    wiou = 1 - (inter + 1) / (union - inter + 1)
    return mean(wbce + wiou)

Inference, as implemented in Test.py:

# Prediction for one image  [code: Test.py#L22-L42]
# [code: utils/dataloader.py#L130-L131] for the ImageNet normalisation
# image: RGB; gt_size: (h, w) of the ground-truth mask

def predict(image, gt_size):
    # Resize to 352 (code default) and normalise  [code: Test.py#L10]
    inputs = normalise(resize(image, 352))
    # Upsample the logit map to the ground-truth size  [code: Test.py#L38]
    logits = upsample_bilinear(model(inputs), size=gt_size)
    # Sigmoid, then min-max normalise the map  [code: Test.py#L39-L40]
    probs = sigmoid(logits)
    return (probs - min(probs)) / (max(probs) - min(probs) + 1e-8)

Hyperparameters

The paper gives two training setups §4.2. The code's defaults match the second one. Values marked "code" are not stated in the paper.

Results

Author-reported. Everything below is copied from the paper's tables; the baseline values in Table 1 are taken by the authors from jha2020real [tab. 1 caption]. FPS is measured on a GeForce RTX 2080 Ti.

Setup A, Kvasir-SEG 880/120 split §4.2 tab. 1: training at 512x512, SGD with learning rate 1e-2 for 100 epochs, with augmentation.

Setup B, PraNet protocol §4.2 tab. 2 tab. 3: 1450 training images from Kvasir-SEG and CVC-ClinicDB, no augmentation, training at 312x312 with Adam (learning rate 1e-4) for 100 epochs, and testing on five datasets. The FPS in Table 2 is at 312x312.

Kvasir-SEG tab. 2:

ClinicDB, ColonDB, ETIS and CVC-T (CVC-T is the test set of EndoScene) tab. 3:

Headline claims §5: HarDNet-MSEG is the only network in the comparison to pass 0.90 mean Dice on Kvasir-SEG, at 0.912 under setup B and 0.904 under setup A. The authors report it as about 1.3 times faster than PraNet and more than twice as fast as the other compared models. The abstract's 0.904 and 86.7 FPS are the setup A figures tab. 1.

Paper vs code

Compared with the released code at aac1c4b5d2fdbc80bc8fe21a5377af2d202a0348. The official repository is the one linked from the abstract.

• Training loss. The paper states no loss §4.2. The code uses the structure loss (weighted BCE plus weighted IoU) on the final map only, and its auxiliary losses are commented out Train.py.
• Optimizer for setup A. The paper states SGD with learning rate 1e-2 §4.2. The code's default is Adam with learning rate 1e-4 Train.py, which matches setup B. SGD is available through the flags, with momentum 0.9 and weight decay 1e-4 Train.py. The paper does not say which configuration produced the setup A numbers.
• Input size. The paper uses 512x512 for setup A and 312x312 for setup B §4.2 tab. 2. The code defaults to 352 for training and for testing Train.py Test.py.
• Epochs. The paper says 100 §4.2. The code loops range(1, opt.epoch), which runs 99 epochs Train.py.
• Learning-rate schedule. The paper states none. The code multiplies the learning rate by 0.1 every 200 epochs, which means a constant rate over 100 epochs Train.py utils.py.
• Multi-scale training. The paper does not mention it. The code runs each batch at 0.75x, 1x and 1.25x of the training size, with one optimizer step per scale Train.py. At the default size of 352, the three sizes are 256, 352 and 448 (derived from the rounding in Train.py#L83).
• Gradient clipping. The paper states none. The code clamps each gradient element to ±0.5, not the gradient norm Train.py utils.py.
• Batch size. Not in the paper. The code defaults to 16 Train.py.
• Augmentation. The paper lists random rotation and horizontal and vertical flips for setup A §4.2. The code applies rotations of up to ±90° and flips with probability 0.5 dataloader.py. It enables augmentation only when the flag is exactly the string True dataloader.py, and it is off by default Train.py.
• Decoder fusion. The paper describes upsampling and element-wise multiplication §3.2.2 fig. 5. The code also concatenates features and applies 3x3 convolutions between the steps HarDMSEG.py.
• Which backbone levels are used. The paper says the shallow features are discarded §3.2 but does not say how many levels. In the code, the backbone outputs are taken at strides 4, 8, 16 and 32 (128, 320, 640 and 1024 channels). The decoder uses only the last three, so the stride-4 output is dropped hardnet_68.py HarDMSEG.py. The strides are derived from the stem and pooling layers in hardnet_68.py.
• ReLU in convolution blocks. The paper does not state this. BasicConv2d defines a ReLU, but its forward pass returns only convolution and batch norm. ReLU is applied only at the end of each RFB block HarDMSEG.py HarDMSEG.py.
• Unused modules. The reverse-attention branches and the layers conv2 to conv6 are created in __init__ but never used in forward HarDMSEG.py HarDMSEG.py. They still appear in model.parameters().
• Pretrained backbone. The paper does not state it §3.1. The code loads a hardnet68.pth file from a hard-coded absolute path on the author's machine hardnet_68.py. The README says to download the file from the Pytorch-HarDNet repository and edit that line README.md.
• Software. The paper states PyTorch 1.6 and CUDA 10.2 §4.4. The README asks for Python 3.6 and PyTorch 1.1 README.md.
• Checkpoint selection. The paper does not state it. The code saves the checkpoint with the highest mean Dice on the folder given by --test_path, which defaults to the Kvasir test folder Train.py Train.py. See Pitfalls.
• Setup A split. The paper cites jha2020real for the 880/120 split §4.2. The code does not create this split. It reads whatever folders are passed in Train.py.

Pitfalls

• Checkpoint selection on the test set. The training script evaluates the mean Dice on the test folder after every epoch and keeps the best checkpoint Train.py. The reported Kvasir numbers may therefore be selected on the test set. The paper does not mention this. A fair reproduction should fix the epoch count in advance or use a separate validation split.
• "Mean Dice" is not one number. The paper gives per-pixel Dice formulas §4.3 but does not say whether the tables average per image, per dataset, or over thresholds. The authors' MATLAB evaluation averages Dice over 256 thresholds (1 to 0 in steps of 1/255) and also reports the best threshold. The per-image map is first min-max normalised main.m main.m main.m. The README says the report uses this MATLAB code README.md. The Python script in eval_Kvasir.py is said to differ from the report by about 0.002 README.md.
• Python Dice in Train.py is not the table metric. The test loop in Train.py computes a soft Dice with smoothing 1 on min-max normalised probabilities. It averages over 100 hard-coded iterations, not over the actual test set size Train.py. Its values are for checkpoint selection, not for comparison with the tables.
• F2 formula. The paper prints F2 as 5pr/(4p+r) without defining p and r §4.3. The usual reading is precision and recall, but the paper does not say so.
• Speed numbers depend on undocumented inputs. The paper gives the input size only for setup B tab. 2. It does not say whether FPS includes the upsampling to ground-truth size or the sigmoid §4.4. The code's default size of 352 is not the 312 that the speed claim in Table 2 uses.
• Multi-scale training costs about three times the compute per batch. Each batch makes three forward and backward passes and three optimizer steps Train.py. Epoch counts and wall-clock times from the paper therefore do not transfer directly to a single-scale run.
• Silent sample dropping. The training loader drops image and mask pairs whose sizes differ, without any warning dataloader.py.
• Hard-coded paths. The pretrained weight path hardnet_68.py, the dataset defaults Train.py and the test folders Test.py point to the authors' machine. Test.py loops over five dataset names, including CVC-300, and expects their folder layout.
• Deprecated calls. Test.py saves maps with scipy.misc.imsave Test.py Test.py, which current SciPy releases no longer provide. Train.py imports torchstat Train.py, which the README does not list.
• Single runs. The paper reports one run per model and gives no variance or seeds §4. Train.py sets no seed Train.py.
• Baselines come from other papers. The setup A caption says the other models' scores are taken from jha2020real tab. 1 caption, so they were not retrained by the authors. The caption does not say whether the baseline FPS values were measured by the authors.
• Two headline numbers. The abstract's 0.904 and 86.7 FPS come from setup A, while 0.912 and 88 FPS come from setup B at 312x312 §5 tab. 2. The two setups differ in input size, optimizer, augmentation and training data, so they should not be treated as one benchmark.

Lineage

• Builds on: HarDNet, the backbone, with its own low-memory-traffic design §3.1. The paper cites it as chao2019hardnet; no card exists yet.
• Builds on: Cascaded Partial Decoder, whose decoder design the paper follows §3.2. The paper cites it as wu2019cascaded; no card exists yet.
• Builds on: Receptive Field Block, used for each decoder input §3.2.1. The paper cites it as liu2018receptive; no card exists yet.
• Trained with the protocols of PraNet §4.2 (setup B). PraNet has no card yet.
• Uses the setup A protocol of jha2020real §4.2.
• Compared against: PraNet, U-Net and its variants, DoubleU-Net, ResUNet and ResUNet++, FCN8, HRNet, PSPNet, DeepLabv3+, SFA, and others tab. 1 tab. 2 tab. 3.
• Superseded by: not known from the content of this wiki.