Spatiotemporal Vision Transformer for Short Time Weather Forecasting

Weather forecasting is a critical research area that may lead to serious consequences, if not done accurately. In order to manage the impact of adverse weather effects short time forecasting is employed, since it can be highly accurate to predict a short time, e.g., hours, into the future. In this work, we tackle the challenge of short time weather forecasting using a novel approach based on a modified UNet based model. Specifically, all convolution-based building blocks were replaced by 3D shifted window transformers in both encoder and decoder branches. Shifted window transformers greatly reduce computational complexity requirements, a major constraint of self-attention, without sacrificing performance. To support the pre-training of the transformer-based model a carefully crafted augmentation scheme was proposed. The model was tested on the IEEE Big Data Weather4cast Competition data, which requires the prediction of 8 hours ahead frames (4 per hour) from an hourly weather product sequence. We show the importance of including other weather products in encouraging spatial generalization, while this may not be optimal for temporal generalization. The model demonstrates highly competitive performance on both the validation and test datasets. The code is available online at https://github.com/bojesomo/Weather4Cast2021-SwinUNet3D