Spectral–Spatial Masked Transformer With Supervised and Contrastive Learning for Hyperspectral Image Classification

Abstract

Recently, due to the powerful capability at modeling the long-range relationships, Transformer-based methods have been widely explored in many research areas including hyperspectral image (HSI) classification. However, because of lots of trainable parameters and the lack of inductive bias, it is difficult to train a Transformer-based HSI classifier, especially when the number of training samples is limited. To address this issue, in this study, spectral-spatial masked Transformer (SS-MTr) is explored for HSI classification, which uses a two-stage training strategy. In the first stage, SS-MTr pre-trains a vanilla Transformer via reconstruction from masked HSI inputs, which embeds the local inductive bias into the Transformer. In the second stage, the well pre-trained Transformer is cooperated with a fully connected layer and then fine-tuned for the HSI classification. Furthermore, in order to incorporate discriminative feature learning into the SS-MTr, three SS-MTr-based methods, including contrastive SS-MTr (C-SS-MTr), supervised SS-MTr (S-SS-MTr), and supervised contrastive SS-MTr (SC-SS-MTr) are proposed by adding extra branches for specific tasks in parallel with the existing reconstruction task. Specifically, the proposed C-SS-MTr adds a contrastive loss which brings instance discriminability. Besides, the proposed S-SS-MTr builds an extra classification branch for embracing inter-class discriminability and intra-class similarity. Moreover, the proposed SC-SS-MTr combines C-SS-MTr and S-SS-MTr for better generalization. The proposed SS-MTr, C-SS-MTr, S-SS-MTr, and SC-SS-MTr are tested on three popular hyperspectral datasets (i.e., Indian Pines, Pavia University, and Houston). The obtained results reveal that the proposed models achieve competitive results compared with the state-of-the-art HSI classification methods. Code is available at https://github.com/mengduanjinghua/SS-MTr.