SmaAt-UNet: Precipitation Nowcasting using a Small Attention-UNet Architecture

In plain words

The paper asks a neural network to look at the last hour of rain radar pictures and draw the rain picture 30 minutes later. Weather services usually predict rain with physics simulations that are slow and need big computers. The authors argue that a smaller model needs less training and computing power, which matters for efficiency and for smaller computers.

The network is a U-shaped image model. It takes the picture, shrinks it step by step to find patterns, then grows it back to full size. Two changes make it smaller and more focused. A small attention step lets it weigh which features matter, and lighter convolutions replace the usual ones. The result has about a quarter of the parameters of the original, and its accuracy stays close to the original's. The predicted pictures are blurry, which is a known side effect of the training loss.

TL;DR

• Task: predict the rain map 30 minutes after the last of 12 five-minute radar maps, pixel by pixel, on a 288 × 288 grid §4.1.
• Model: U-Net with CBAM attention after the entry block and after each of the four encoder stages, and depthwise-separable convolutions with two kernels per layer in place of ordinary ones §3.1, §3.2, fig. 5.
• Size: 4,111,389 parameters against 17,272,577 for U-Net tab. 1.
• Result (author-reported, NL-50 test set, 2019): MSE 0.0122, the same as U-Net; NMSE 416.10 against U-Net's 416.38; slightly behind U-Net on most other scores tab. 4.
• Watch for: the training loss is a per-image sum, not the paper's mean; the NL-50 filter in the released code checks a different frame than the target; the dataset is not public. See Paper vs code and Pitfalls.

Problem

• Numerical weather prediction is accurate over hours to days but is computationally expensive and unsuited to nowcasts from minutes up to six hours ahead §1.
• Data-driven nowcasting models can use the latest observations, but at very short horizons the persistence baseline (predict that the last input map stays unchanged) is hard to beat §1, §2.
• On the NL-50 test set every model in the paper beats persistence by a large margin: persistence MSE 0.0248 against 0.0122 for the best models tab. 4 §5.1.
• Assumptions: the precipitation maps come from two C-band Doppler radars in De Bilt and Den Helder, rain-gauge adjusted, at 5-minute intervals, covering 2016–2019 §4.1. The dataset is available only by request to the authors §4.1.
• Rain is rare in the raw data. Training on all samples biases the network towards predicting zero, so the authors train on samples whose target map has at least 50% rainy pixels §4.1.

Method

Overview

The input is the 12 maps stacked along the channel axis, shape (B, 12, 288, 288). The output is one map (B, 1, 288, 288) with the predicted rain value for each pixel §4.1 §3.1. The network has the four encoder–decoder modules of the original U-Net §3.1.

The modifications are:

1. CBAM attention (convolutional block attention module) after the first double convolution and after each encoder stage. A CBAM applies channel attention, then spatial attention §3.1 (Woo et al., 2018, cited in §2). The attended features feed only the skip connections into the decoder. The features passed to the next encoder stage are not attended, so the encoder path is unchanged §3.1.
2. Depthwise-separable convolutions in place of the regular convolutions of U-Net. The paper says "all convolutions" are replaced, but the CBAM spatial convolution stays regular §3.1. The final 1 × 1 output convolution is also regular in the code SmaAt_UNet.py unet_parts.py.

Channel counts are not stated in the paper. The code uses 64, 128, 256, 512 for the encoder and 512 at the bottleneck with bilinear upsampling SmaAt_UNet.py.

The bottleneck feature map is 18 × 18 §5.1 fig. 5. Tensor sizes below follow from 288 × 288 input and four 2 × 2 poolings.

# SmaAt-UNet forward pass  [§3.1] [fig. 1]
# [code: models/SmaAt_UNet.py#L41-L57]
# Input:  rain maps (B, 12, 288, 288), normalized
# Output: predicted map (B, 1, 288, 288), normalized

def forward(maps):
    # Entry block with two depthwise-separable convs  [§3.1]
    level0 = double_dsc_conv(maps)                  # (B, 64, 288, 288)
    # Attention on the skip path only  [§3.1]
    skip0 = cbam0(level0)
    # Encoder: pool, then double DSC conv on the un-attended features  [§3.1]
    level1 = double_dsc_conv(max_pool(level0))      # (B, 128, 144, 144)
    skip1 = cbam1(level1)
    level2 = double_dsc_conv(max_pool(level1))      # (B, 256, 72, 72)
    skip2 = cbam2(level2)
    level3 = double_dsc_conv(max_pool(level2))      # (B, 512, 36, 36)
    skip3 = cbam3(level3)
    # Bottleneck, with the fifth attention layer  [§3.1] [fig. 5]
    bottleneck = double_dsc_conv(max_pool(level3))  # (B, 512, 18, 18)
    attended = cbam4(bottleneck)
    # Decoder: upsample x2, concatenate skip, double DSC conv  [§3.1]
    x = decode(attended, skip3)                     # (B, 256, 36, 36)
    x = decode(x, skip2)                            # (B, 128, 72, 72)
    x = decode(x, skip1)                            # (B, 64, 144, 144)
    x = decode(x, skip0)                            # (B, 64, 288, 288)
    # Final 1 x 1 convolution to a single channel  [§3.1]
    return conv1x1(x)                               # (B, 1, 288, 288)

Note: the lines above are for the algorithm's flow. The skip connection is concatenated before the upsampled tensor in the code unet_parts_depthwise_separable.py. The order does not change the result but changes the weight layout.

Depthwise-separable convolution

A regular 3 × 3 convolution is split into a per-channel depthwise convolution followed by a 1 × 1 pointwise convolution §2. In SmaAt-UNet the depthwise step produces kernels_per_layer = 2 outputs per input channel §3.2 layers.py.

# Depthwise-separable 3 x 3 conv, kernels_per_layer = 2  [§3.2]
# [code: models/layers.py#L34-L50]
# Each input channel gets 2 depthwise filters; the 1 x 1 conv mixes channels

def dsc(x):
    # Depthwise: 2 outputs per input channel, groups = channels_in  [§3.2]
    x = depthwise_conv3x3(x, channels_out = 2 * channels_in)
    # Pointwise: 1 x 1 conv to the block's output width  [§2]
    return conv1x1(x, channels_out = channels_out)

Each DSC double-convolution block is DSC, batch norm, ReLU, DSC, batch norm, ReLU unet_parts_depthwise_separable.py. The paper describes the block only as a "double convolution" §3.1. Batch norm and ReLU are not described in the text.

CBAM

# CBAM: channel attention, then spatial attention  [§3.1]
# [code: models/layers.py#L90-L141]
# Shared MLP C -> C/16 -> C with ReLU (reduction 16 is code-only)

def cbam(x):
    # Channel attention: average- and max-pooled channel vectors  [§3.1]
    avg = global_avg_pool(x)                    # (B, C, 1, 1)
    mx = global_max_pool(x)                     # (B, C, 1, 1)
    gate = sigmoid(mlp(avg) + mlp(mx))          # (B, C, 1, 1)
    x = x * gate
    # Spatial attention: 7 x 7 conv on channel-wise mean and max  [§3.1]
    pooled = concat(mean_over_channels(x), max_over_channels(x))
    gate = sigmoid(batch_norm(conv7x7(pooled)))  # (B, 1, H, W)
    return x * gate

The paper gives neither the reduction ratio nor the spatial kernel size. The code uses 16 regression_lightning.py and 7 layers.py.

Training loss

The paper's loss is the MSE of Eq. 1 eq. 1 §3.4. The code's training loss sums squared errors over every pixel of the batch and divides by the batch size, so it is a per-image sum regression_lightning.py. See Paper vs code.

# Training loss as implemented  [code: models/regression_lightning.py#L57-L65]
# Paper Eq. 1 normalises by the number of samples  [eq. 1]

def loss(predicted, target):
    # Squared error for every pixel in the batch  [eq. 1]
    squared = (predicted - target) ** 2
    # Divide by batch size: sum over pixels, then mean over images
    return sum(squared) / batch_size

Evaluation

Scores are computed on the test maps after converting back to physical units:

# Test-time scoring at 0.5 mm/h  [§3.4] [code: metric/precipitation_metrics.py#L69-L85]
# Normalized values are scaled to mm per 5 minutes  [§4.1]

def score(predicted, target):
    # Undo max normalisation; the value 47.83 is code-only  [§4.1]
    # [code: metric/precipitation_metrics.py#L23]
    predicted = predicted * 47.83
    target = target * 47.83
    # Convert mm per 5 min to mm per hour, times 12
    # [code: metric/precipitation_metrics.py#L80-L81]
    predicted_hourly = predicted * 12
    target_hourly = target * 12
    # Rain mask at 0.5 mm/h  [§3.4]
    rain_predicted = predicted_hourly > 0.5
    rain_target = target_hourly > 0.5
    # Confusion counts for the rain / no-rain decision
    return count_tp_fp_tn_fn(rain_predicted, rain_target)

From these counts:

CSI = \frac{TP}{TP + FN + FP} eq. 2

FAR = \frac{FP}{TP + FP} eq. 3

HSS = \frac{TP \times TN - FN \times FP}{(TP+FN)(FN+TN)+(TP+FP)(FP+TN)} eq. 4

Precision, recall, accuracy and F1 are the standard confusion-matrix definitions. The paper does not write them out precipitation_metrics.py. Recall is also called the probability of detection §3.4.

The MSE of Eq. 1 is

MSE = \frac{\sum_{i=1}^{n} (y_i - \hat{y}_i)^2}{n} eq. 1

where n is the number of samples, y_i the ground truth and \hat{y}_i the prediction §3.4.

Hyperparameters

Radar characteristics tab. 2:

Dataset sizes tab. 3:

Results

Author-reported. Setup: radar precipitation maps (Netherlands), SmaAt-UNet and three comparison models, test on 2019 §4.1 §5.1. The persistence baseline copies the last input map. Bold marks the best value in each column tab. 4. The paper gives no seeds, variance or repeated runs §5.

NL-50 test set, scores on rain above 0.5 mm/h tab. 4:

The text says SmaAt-UNet is second-best in six of the seven scores §5.1. Its NL-50 MSE ties U-Net. Its NMSE is lower, 416.10 against 416.38, by a margin the paper does not comment on.

NL-20 test set, the generalisation check tab. 5. The table's model-size entry for SmaAt-UNet is written "0.24s" in the source; read it as 0.24x:

The paper calls NMSE the fairer comparison on NL-20, because many pixels near zero keep the raw MSE low §5.1. On NL-20 U-Net is slightly better than SmaAt-UNet on NMSE §5.1.

Cloud-cover dataset (binary cloud masks, France, 256 × 256) tab. 6. The persistence row is commented out in the source, and no cloud-cover code is in the repository:

The paper says the four models are "comparable" on this dataset. The largest MSE gap is 0.0019 §5.2.

Paper vs code

Repository: https://github.com/HansBambel/SmaAt-UNet, linked from the paper's footnote in §3.2. Compared at commit 20b99b194cd2a36f76efac1b6b35871db1ce4616, the default branch as cloned.

Differences:

• Early stopping wording. The paper stops training once validation loss has failed to rise for 15 epochs §3.3. The code uses PyTorch Lightning's EarlyStopping in min mode with patience 15 on val_loss train_precip_lightning.py. The counter resets on any new minimum. This is nearly the same rule, but the wording differs.
• Loss normalisation. Eq. 1 divides by the number of samples n eq. 1. The code sums over all pixels in the batch and divides by the batch size regression_lightning.py, so it is a per-image sum of squared errors. The two differ by a scale factor. Adam is largely insensitive to a constant loss scale, but this is an inference, not a test.
• Which MSE is reported. The paper's table column "MSE" does not say which quantity it is. The README says it is mse_pixel [README]. In code that is the denormalised sum of squared errors divided by the number of pixels, i.e. a per-pixel mean in mm per 5 min precipitation_metrics.py. The table's units are not stated.
• NMSE is not computed in the repository. The paper defines NMSE as MSE divided by the average pixel value §5.1. The average value is not stated, and no code in the repository computes NMSE.
• NL-50 filter checks a different frame than the target. The builder keeps a window when frame i, the first frame after the 18-frame window, has at least the required share of rainy pixels create_datasets.py. The stored window is images[i-18:i], and its last frame images[i-1] is the target create_datasets.py dataset_precip.py. The paper states that target images contain at least 50% rain §4.1. This is an inference from the code.
• Output convolution. The paper says all U-Net convolutions are replaced by depthwise-separable ones §3.1. The final 1 × 1 output convolution is a regular convolution SmaAt_UNet.py unet_parts.py.
• Code not in the paper. Batch size 16, CBAM reduction ratio 16, CBAM kernel 7, BatchNorm and ReLU in each block, the 64–512 channel counts, and the 47.83 scale are in the code only train_precip_lightning.py regression_lightning.py layers.py SmaAt_UNet.py.
• Repository branch. The README says the code behind the paper is on branch snapshot-paper, and that master has since upgraded packages and was refactored, so results may not reproduce exactly [README]. The snapshot-paper branch was not inspected. The code above is master.
• Persistence baseline. Code returns the last input frame unchanged regression_lightning.py, matching the paper's description §2.
• Scoring formulas. CSI, FAR and HSS in the code match Eqs. 2–4 precipitation_metrics.py. The 0.5 mm/h threshold is applied after conversion to mm per hour precipitation_metrics.py, as the paper states §3.4.
• Cloud-cover experiments. No code for them is in the repository. The README mentions no cloud pipeline [README].

Pitfalls

• Blurry output. The paper attributes the blur to the MSE loss, which averages over possible futures §5.1. Expect blurry maps from this training setup.
• Scale of the loss. Eq. 1 and the code's loss differ by a normalisation factor (see Paper vs code). If you replicate the code, keep its loss; if you follow the paper, the reported numbers may not match.
• Off-by-one in the NL-50 filter. Fix the frame index if you build your own NL-50 set, or reproduce the code's behaviour knowingly. Both choices change the training distribution.
• Dataset access. The dataset is available only by request to the authors §4.1 [README]. The raw-to-288 preprocessing (cutout of the radar range and centre crop fig. 2) is not in the repository. Normalisation by 47.83 is read from code.
• Unstated values. Batch size, CBAM reduction ratio and kernel size, channel counts and Adam's betas and eps are not in the paper. Use the code's values, and say so when you do.
• Units in the MSE column. The paper does not state the units. Treat the table values as mm per 5 min squared, per the code's mse_pixel.
• NMSE is not reproducible from the repository. The denominator is not given.
• Persistence is strong at short horizons. The paper beats it by a large margin on NL-50 §5.1, but persistence remains the baseline any implementer should beat first.
• Generalisation. The model trains on NL-50 only. On NL-20 the gap to U-Net is small, but the paper notes that training on rain-heavy samples may bias the model toward over-predicting rain §4.1 §5.1.
• Single run per model. Results come from one selected checkpoint per model, with no seeds or variance reported §5.
• Cloud-cover details. Output format, loss adaptation and the use of the six output frames for this dataset are not stated §4.2.
• Licence of the code. The repository has no LICENSE file. Its code is referenced here by path and line only; no code is quoted.

Lineage

• Builds on: U-Net (Ronneberger et al., 2015) §3.1; CBAM (Woo et al., 2018) §2; depthwise-separable convolutions as in MobileNets (Howard et al., 2017) §2. These are named, not linked, because they have no wiki cards.
• Compared against: persistence, U-Net, U-Net with CBAM, U-Net with DSC §5.1. The related work also names ConvLSTM and PredRNN, which are not compared here §2.
• Superseded by: not stated in the paper.