TernausNet: U-Net with VGG11 Encoder Pre-Trained on ImageNet for Image Segmentation

In plain words

Image segmentation means labelling every pixel of a picture. Here the labels are "building" and "not building" in aerial photos. The model is a U-Net: one half shrinks the image into features, the other half grows them back into a full-size mask, and shortcut links carry fine detail across from the first half to the second.

The paper swaps the first half for a small image classifier, VGG11, whose weights were learned on ImageNet. Starting from those learned features, instead of random numbers, makes the model converge faster and reach a higher score on held-out images. The authors also show that a model pre-trained on a different, large segmentation competition dataset (Carvana) does about as well.

The paper is short. It reports one dataset, one run per setting and a small validation set, so the gains should be read as a single experiment rather than a broad result.

TL;DR

• A U-Net whose encoder is the convolutional part of VGG11, initialised from ImageNet weights, segments one class per pixel. Training uses binary cross-entropy minus the log of a soft Jaccard (IoU) term §2 eq. 4.
• On the Inria aerial dataset, validation IoU after 100 epochs is 0.593 with random (LeCun-uniform) initialisation, 0.686 with the pretrained encoder and 0.687 with a fully pretrained network §3. These are author-reported numbers from 30 validation images.
• The official repository (MIT) contains only the network definition. The training loop, loss and pretrained weights are not in the cloned commit.
• Paper and code disagree in a few details: the encoder has 7 convolutions in the text but 8 in the code, and the skip connection is concatenated at a different point (see Paper vs code).

Problem

Dense segmentation needs a mask for every training image. Masks are costly to make, so most segmentation datasets hold at most thousands of images. Training a U-Net from random weights on such a set is hard to avoid overfitting, and the usual remedy is to start from weights learned on a large classification dataset such as ImageNet §1 §2.

The paper assumes binary segmentation, one class against background §1 §3. It also assumes input sides are divisible by 32, because the network downsamples five times by 2 §2.

Method

The network is a U-Net with a VGG-style encoder §2 fig. 1.

1. Encoder. The fully connected layers of VGG11 are removed. The encoder keeps the convolutional stages, each convolution followed by ReLU, with 2x2 max-pooling between stages. Channels go 64, 128, 256, 512 and stay at 512 §2 fig. 2.
2. Bottleneck. The paper describes one 512-channel convolution between encoder and decoder §2. The code's bottleneck differs (see Paper vs code).
3. Decoder. Five upsampling steps, each a transposed convolution that doubles the spatial size and halves the channels, pair with the five pooling steps §2 fig. 1. Each step concatenates the matching encoder feature map §1 §2.
4. Output. A 1x1 convolution produces one channel per pixel, a probability once a sigmoid is applied. The paper's ŷ is that probability §3 eq. 2. The network itself returns logits models.py.
5. Initialisation. Three settings are compared §3: (a) LeCun-uniform for all layers, with bound L = sqrt(1/f_in); (b) VGG11 encoder pretrained on ImageNet, decoder LeCun-uniform; (c) the full network pretrained on Carvana.
6. Training. Adam with learning rate 0.001, 100 epochs, random 768x768 crops of the training images §3.
7. Loss. The loss is L = H − log J eq. 4. Here H is pixel-wise binary cross-entropy eq. 3, and J is the mean over pixels of a soft Jaccard term, y·ŷ / (y + ŷ − y·ŷ) eq. 2. The Jaccard index J(A, B) is defined on sets eq. 1.
8. Inference. The output probability map is thresholded at 0.3, a value the paper says can be found on the validation set. Pixels above it become 1, the rest 0 §2.

Pseudocode of the forward pass, with the bottleneck and the decoder written as the code does them:

# TernausNet forward pass  [§2] [fig. 1]
# Input:  image (B, 3, H, W); H and W divisible by 32  [§2]
# Output: logits (B, 1, H, W); no sigmoid  [code: ternausnet/models.py#L108]

def forward(image):
    # Stage 1, full resolution, 64 channels: one conv and ReLU  [§2] [fig. 2]
    # [code: ternausnet/models.py#L92]
    skip1 = relu(conv1(image))                           # (B, 64, H, W)

    # Stage 2, 128 channels: max-pool, then one conv and ReLU  [§2] [fig. 2]
    # [code: ternausnet/models.py#L93]
    skip2 = relu(conv2(max_pool(skip1)))                 # (B, 128, H/2, W/2)

    # Stage 3, 256 channels: max-pool, then two convs with ReLU  [§2] [fig. 2]
    # [code: ternausnet/models.py#L94-L95]
    skip3 = relu(conv3(relu(conv3s(max_pool(skip2)))))   # (B, 256, H/4, W/4)

    # Stage 4, 512 channels: max-pool, then two convs with ReLU  [§2] [fig. 2]
    # [code: ternausnet/models.py#L96-L97]
    skip4 = relu(conv4(relu(conv4s(max_pool(skip3)))))   # (B, 512, H/8, W/8)

    # Stage 5, 512 channels: max-pool, then two convs with ReLU  [§2] [fig. 2]
    # [code: ternausnet/models.py#L98-L99]
    skip5 = relu(conv5(relu(conv5s(max_pool(skip4)))))   # (B, 512, H/16, W/16)

    # Bottleneck: 512 to 512 conv, then 2x upsampling to 256 channels  [§2]
    # [code: ternausnet/models.py#L72-L74] [code: ternausnet/models.py#L101]
    center = decoder_block(max_pool(skip5))             # (B, 256, H/16, W/16)

    # Decoder block: 3x3 conv and ReLU, then 2x transposed conv and ReLU  [§2]
    # [code: ternausnet/models.py#L24-L44]
    # The skip is concatenated before the block, not after it  [code: L103-L107]
    dec5 = decoder_block(concat(center, skip5))         # (B, 256, H/8, W/8)
    dec4 = decoder_block(concat(dec5, skip4))           # (B, 128, H/4, W/4)
    dec3 = decoder_block(concat(dec4, skip3))           # (B, 64, H/2, W/2)
    dec2 = decoder_block(concat(dec3, skip2))           # (B, 32, H, W)
    dec1 = conv_relu(concat(dec2, skip1))               # (B, 32, H, W)

    # Final 1x1 convolution to one channel  [§3]
    # [code: ternausnet/models.py#L89]
    return conv1x1(dec1)                                # (B, 1, H, W)

Loss, as stated in the paper (pseudocode, not in the repository):

# Training loss L = H - log J  [§3] [eq. 4]
# Input:  logits (B, 1, H, W); targets (B, 1, H, W) with values 0 or 1
# Output: scalar loss

def loss(logits, targets):
    # Predicted probability per pixel, the paper's y-hat  [§3] [eq. 2]
    probs = sigmoid(logits)
    # Pixel-wise binary cross-entropy, averaged over pixels  [§3] [eq. 3]
    bce = -mean(targets * log(probs) + (1 - targets) * log(1 - probs))
    # Soft Jaccard per pixel, y * p / (y + p - y * p), averaged  [§3] [eq. 2]
    intersection = targets * probs
    jaccard = mean(intersection / (targets + probs - intersection))
    # Combined loss  [§3] [eq. 4]
    return bce - log(jaccard)

Hyperparameters

Results

Author-reported. Setup: Inria Aerial Image Labeling, 150 training images and 30 validation images, U-Net with VGG11 encoder, 100 epochs, Jaccard index (IoU) on the validation set §3. Baseline: the same network with LeCun-uniform initialisation. The table below is compiled from the text; the paper has no results table.

The paper also quotes two reference values for the same validation split, as reported by other works: a validation IoU of about 0.647 §3 and a best validation IoU of about 0.73 §3. The paper says its results can be improved with hyperparameter search and standard pre- and post-processing, and gives no numbers for either §3.

Curves show that both pretrained models converge faster than the random one, and the authors say the steady-state value seems higher for them §3 fig. 3. Masks for the three settings are shown qualitatively fig. 4; no metric is attached to them.

Paper vs code

1. Encoder depth. The text says VGG11 has seven convolutional layers §2. The code takes eight convolutional modules from torchvision's VGG11 features (indices 0, 3, 6, 8, 11, 13, 16, 18) models.py. Implement the eight the code uses; the text and code disagree, and the code is what the authors released.
2. Skip concatenation. The text describes upsampling first, then concatenation, then a convolution §2. The code concatenates the encoder map first and feeds the result into the decoder block, which upsamples it models.py models.py.
3. Bottleneck. The text describes a single 512-channel convolution that replaces the fully connected layers §2. The code's bottleneck is a 512-to-512 convolution followed by a 2x transposed convolution to 256 channels, applied to the fifth pooled map models.py models.py.
4. Output. The paper's ŷ is a probability §3 eq. 2. The model returns logits and applies no sigmoid models.py. The repository has no loss or inference code, so the sigmoid must be added by the caller.
5. Weight initialisation. The paper says all decoder layers use LeCun-uniform and that this is PyTorch's convolution default §3. The code sets no initialisation and relies on the framework models.py models.py. Whether the framework default equals the stated formula is not checked here.
6. Pretrained encoder. The code builds the encoder with models.vgg11(pretrained=...) and documents True as "encoder is pre-trained with VGG11" models.py. The repository does not state which weights are loaded or which torchvision version is needed.
7. Released weights. The abstract says the pretrained weights are published [abstract]. The cloned commit contains no weight files. The README tells users to pip install ternausnet README.md; whether that package carries the Carvana weights was not checked.
8. Missing parts. Training loop, loss, data pipeline, threshold and evaluation are not in the repository. The Python source is ternausnet/models.py only models.py. The paper's training details must be taken from the text.
9. Extra model. The same file defines UNet16, with a VGG16 encoder and optional deconvolution decoder models.py. The paper names VGG16 only as future work §4.
10. License. The conclusion says the code is MIT §4. The LICENSE file is MIT, with copyright held by the two authors LICENSE.

Pitfalls

• Citation numbering. § numbers here are the four \section commands in order (Introduction §1, Network Architecture §2, Results §3, Conclusion §4). The typeset PDF prints them as I to IV. Equation numbers are the four numbered equation environments in order, eq. 1 to eq. 4. Figures are numbered in order of appearance: fig. 1 U-Net, fig. 2 VGG11, fig. 3 learning curves, fig. 4 masks.
• Sigmoid placement. Eq. 3 needs ŷ in (0, 1), and the 0.3 threshold applies to probabilities §2. The model returns logits models.py, so apply the sigmoid before the loss and before thresholding.
• 0/0 in the Jaccard term. For a pixel with y = 0 and ŷ = 0, eq. 2 has denominator 0 eq. 2. Implementations need an epsilon or a guard, which the paper does not state.
• log(0). Both the BCE term and −log J are undefined when a probability or the mean Jaccard is exactly 0 eq. 3 eq. 4. Clamping is not stated.
• Reduction not stated. The paper does not say whether the Jaccard mean and the BCE mean run over the whole batch or per image eq. 2 eq. 3. The pseudocode above averages over all pixels in the batch; that is an inference from the equations, not a stated choice.
• IoU reporting not stated. The paper does not say whether IoU is computed per image and averaged or over the whole validation set §3 eq. 1.
• Small validation set. Results use 30 validation images, one run per setting, no seeds and no variance §3. The paper reports no test-set result.
• Single dataset. The paper claims the 0.3 threshold is "pretty universal" for this loss and many datasets §2, but the results cover only the Inria aerial dataset §3.
• Carvana model. The third setting (0.687) uses a network pretrained on Carvana. The paper does not give the size, epochs or loss of that pretraining §3.
• Garbled sentence. The results paragraph contains a broken sentence about learning curves ("validation learning curves in Our results ...") §3. The surrounding claims about steady-state values are qualitative §3 fig. 3.
• Torchvision dependency. The pretrained flag of torchvision model constructors is version-dependent. Check that the installed torchvision accepts it before relying on the code path models.py.
• Divisibility. The network needs input sides divisible by 32 §2. The code has no check. Each pooling rounds the size down, so other sizes can give mismatched shapes at the concatenations (inferred from the code, not tested).
• Padding and output size. The decoder's transposed convolution uses output_padding 1 with kernel 3 and stride 2 models.py. Changing the kernel or padding alters the output size.

Lineage

• Builds on: U-Net, arXiv:1505.04597 §1 §2; VGG, arXiv:1409.1556, for the encoder §2; Adam, arXiv:1412.6980 §3. No cards for these exist yet.
• Compared against: the same network with LeCun-uniform initialisation §3. Two reference IoU values from other works are quoted §3.
• Future work named by the paper: VGG16 and ResNet encoders §4. The repository's UNet16 covers the VGG16 case models.py.
• Superseded by: not known from content already in this wiki.