VTR: An Optimized Vision Transformer for SAR ATR Acceleration on FPGA

Deep neural networks have gained high interest, showcasing impressive results across various problem domains. The state-of-the-art works for SAR ATR applications involve either convolutional neural networks (CNNs) or Graph Neural Networks (GNNs)[24]. Morgan [25] is the first work to propose a deep CNN for SAR ATR. Recently, Zhang et al.[7] proposed a novel architecture based on GNNs that exceeded classification accuracy of 99% on the MSTAR dataset. In contrast to CNNs, the proposed GNN exploits the data sparsity in SAR images to reduce computation complexity. Although ViTs have emerged as state-of-the-art models for computer vision tasks, they perform poorly on small datasets due to severe overfitting. Hence, training a ViT model for SAR ATR applications without pre-training is challenging. Li et al.[26] proposed a pre-processing technique to improve the accuracy of the ViT model on SAR image classification. Recently, Shifted Patch Tokenization (SPT) and Locality Self-Attention (LSA) techniques have been proposed to improve the accuracy of ViTs on small datasets[20]. In this work, we propose VTR model for SAR ATR, which incorporates the SPT and LSA modules into a vanilla ViT model. This enables training VTR on small SAR datasets without pretraining.

ViT[9] has achieved significant advancements in computer vision tasks like image classification and object detection. Its core components include Multi-Headed Self-Attention (MSA) and Multi-Layer Perceptron (MLP) blocks. First, the input image undergoes partitioning into non-overlapping patches. Then, the flattened patches are embedded along with the class token and positional information. Subsequently, this processed data is fed into the transformer encoder.

Multi-Headed Self-Attention: In the self attention block, the input embeddings are linearly projected to query, key, and value vectors. The query and key vectors are then utilized to obtain a scaled dot product. A single self-attention operation is denoted as one ”head”. The self-attention function is defined as,

where $\bm{Q}$, $\bm{K}$ and $\bm{V}$ represent the query, key, and value matrix, respectively. $d_{k}$ is the dimension of embeddings in $\bm{K}$. Several such heads are concatenated and projected via a linear layer to compute the final MSA output.

Multi-Layer Perceptron: The output of the MSA layer is fed into the MLP block. It consists of two linear layers with an activation function. The MLP layer is formulated as,

where $\bm{W}_{1}$ and $\bm{b}_{1}$ represent the hidden layer weights and bias, respectively. $\bm{W}_{2}$ and $\bm{b}_{2}$ are the output layer weights and bias respectively. GeLU is the activation function.

In the standard ViT model, the receptive field of the input embeddings (input tokens) remains fixed and cannot be adjusted. That is, the tokenization of the standard ViT is similar to the operation of a non-overlapping convolutional layer[20]. As a result, these tokens have a small receptive field leading to lower local inductive bias. To this end, SPT leverages spatial information by shifting the input image in the four or eight cardinal directions. Additionally, since both the query and key are linearly projected from the same input tokens, the self-token relations tend to exhibit larger magnitudes compared to inter-token relations. Consequently, the softmax function assigns relatively higher scores to self-token relations and smaller scores to inter-token relations. Hence, the attention of standard ViT tends to be similar to each other regardless of inter-token relations. Moreover, the scaling factor $\sqrt{d_{k}}$ can cause smoothing of the attention score distribution[27]. These issues can be effectively mitigated by incorporating LSA into the attention layer. LSA helps in excluding self-token relations and applies a learnable temperature scale to the softmax function.

Our model first transforms the input image using SPT. Subsequently, both the original input image and the transformed images (shown in Figure 2) are concatenated and partitioned into non-overlapping patches. The flattened patches undergo layer normalization and linear projection to obtain patch embeddings which are concatenated with a learnable class embedding and then added to learnable positional embeddings to yield the final input tokens. This is then fed into the transformer encoder where the input is passed through several layers of multi-headed self-attention and multi-layer perceptron networks. In the multi-headed self-attention layers, our model incorporates the LSA module. The LSA module applies a learnable temperature scaling to the softmax function to sharpen the distribution of attention scores. Layer normalization is applied before every attention and MLP block while residual connections are applied after each block. Finally, the output of the transformer encoder is passed through an MLP head to generate the classification result.

In the SPT module, each input image, $\bm{x}\in\mathbb{R}^{H\times W\times C}$, is shifted by $2$ pixels in the four diagonal directions: left-up, right-up, left-down, and right-down. These shifted images are cropped to the same size as the original input image and concatenated with it (Figure 2). The resulting set of images are then divided into $N$ non-overlapping patches and flattened into a sequence of vectors, $\bm{x}^{i}\in\mathbb{R}^{P^{2}.C.(N_{s}+1)}$. $H,W,$ and $C$ represent the height, width and the number of channels of the input image respectively. In our case, $C=1$. $\bm{x}^{i}$ is the $i^{th}$ flattened vector. $P$ indicates the size of the patch while $N_{s}$ represents the total number of shifted images. The patch embeddings (flattened vectors) are linearly transformed into a hidden dimension, $d_{s}$. A class token is concatenated to the linearly transformed patch embeddings and positional embeddings are added. The output of this SPT module is fed to the transformer encoder.

The main components of the LSA module include diagonal masking and learnable temperature scaling. In the multi-headed self-attention layer, two linear projections of the input are generated: the query, $\bm{Q}=\bm{X}\bm{W}_{q}$, and the key, $\bm{K}=\bm{X}\bm{W}_{k}$. As formulated in eq.1, the softmax function is applied on a scaled dot product of these linear projections. The diagonal of the resultant matrix of the dot product, $\bm{Q}\bm{K}^{T}$, represents the self-token relations while the off-diagonal elements represent the inter-token relations. Since the linear projections are from the same input tokens, the self-token relations tend to be larger resulting in higher scores for self-token relations. To prevent this the LSA module forces $-\infty$ on the diagonal elements. This effectively prevents the attention from being focused on its own tokens. The learnable temperature scaling technique incorporated in the LSA module, allows the ViT to decide the softmax temperature by itself while training.

Figure 3 describes the data layout. In order to compute the matrix product $\bm{A}\bm{B}$, $\bm{A}$ being the left matrix and $\bm{B}$ the right matrix, the left matrix $\bm{A}$ is stored in a row-major format and the right matrix $\bm{B}$ is stored in a column-major format.

The matrices are divided into blocks of size $b\times b$. The matrices are stored in a block-contiguous fashion, either in row-major format (left matrix) or column-major format (right matrix). The elements within each block are stored in a contiguous fashion.

In this Section, we describe the overall compute flow for the accelerator in Section 5.2 to compute the output for an input $\bm{X}$. The input $\bm{X}$ is shifted and tokenized (Section 3). The compute flow within each VTR encoder layer can be partitioned into two flows: (i) Multi-Headed Self-Attention (MSA) compute (ii) Multi-Layer Perceptron (MLP) compute. Prior to describing the MSA and MLP compute, we first describe the generation of learned embeddings from the shifted tokenized input, $\bm{X}$.

We implement the VTR model using PyTorch 2.0.1[28] and utilize an NVIDIA RTX A6000 GPU with CUDA 11.8 for training the model. Several model variants are trained to comprehensively evaluate the VTR model’s robustness and generalization across different small-sized SAR datasets. Specifically, we vary the following hyper-parameters:

• •

  Patch Size: We explore different patch sizes for partitioning the input image, to assess its impact on model performance.

• •

  Hidden Dimension: The dimensionality of the hidden layers within the transformer encoder is varied to observe its effects on model learning.

• •

  Depth: In our experiments, we explore varying depths of layers in the transformer encoder to assess the model’s sensitivity to depth.

• •

  Number of Heads: We use two settings of attention heads in the multi-headed self-attention mechanism to analyze its influence on model behavior.

We directly train the VTR model on three distinct SAR datasets: MSTAR [21], SynthWakeSAR [22], and GBSAR [23]. The three datasets represent a diverse range of SAR imaging scenarios and characteristics. The Adam optimization algorithm [29] is utilized for training, with an initial learning rate of $0.001$. Subsequently, the step learning rate optimizer is applied, with a step size of $20$ and a gamma value of $0.5$. The model is trained for $200$ epochs.

The HPPU and ECU described in Section 5 are implemented on a state-of-the-art FPGA platform, Xilinx Alveo U250. The main hyper-parameters associated with the proposed accelerator, $p_{h},p_{t},p_{c}$ and $p_{pe}$ are selected as $4,12,2$ and $8$, respectively. $p_{h}$ of $4$ allows the accelerator to fully utilize the $4$ SLRs (Super Logic Regions) on the FPGA. $p_{t}$ was selected based on the nominal input token blocks (associated with the input image). $p_{c}$ of $2$ allows for both the PE columns within each HCU (in each SLR) to concurrently access the FPGA BRAM/URAM using its dual port memory. $p_{pe}$ of $8$ is selected to support any block size $b$ which is a multiple of $8$ (such as $8,16,32$, etc.). The accelerator is designed using Xilinx High Level Synthesis (HLS) and synthesized using Xilinx Vitis v2022.2 with an achieved frequency of $300$ MHz.

MSTAR: The original Moving and Stationary Target Acquisition and Recognition (MSTAR) dataset consists of 5172 SAR images. It contains 10 categories of ground vehicles, with 2747 images in the training set and 2425 images in the testing set. Each image is of $128\times 128$ pixels. For our experiments, we use an augmented MSTAR dataset which consists 27,000 SAR images in the training set and 2425 images in the testing set. The size of each image in this dataset is $88\times 88$. Hence, the images in this dataset are generated by cropping the images in the original MSTAR dataset in various directions.
SynthWakeSAR: SynthwakeSAR is a synthetic SAR imagery dataset comprising 92,160 images of 10 different real vessel models. There are 73,728 images in the training set and 18,432 images in the test set. Each image is of $128\times 128$ pixels. The images of each vessel contain the ship’s wakes.
GBSAR: Ground Based SAR (GBSAR) dataset consists of 6434 raw radar images of which 5147 images are in the training set and 1287 images are in the test set. This dataset captures 7 different ceramic cups with rubber objects. Each raw SAR image has a size of $88\times 88$.

We evaluate the performance of our accelerator for the SAR ATR applications against state-of-the-art CPU and GPU platforms, with the specifications of the platforms detailed in Table 1. Table 1 also contains specifications associated with the FPGA platform utilized for the hardware accelerator.

We evaluate our model using the following metrics: classification accuracy, total MACs (multiply and accumulate operations), model parameters, inference latency and throughput. Classification accuracy is measured as the ratio of images correctly classified by the model. MACs and model parameters characterize the computational complexity associated with a given model; MACs measures the total multiply-accumulate operations that a model performs to compute the output, model parameters are the total learnable parameters in a model. The inference latency is defined as the end-to-end latency (time) required for a model to compute the output given an input image, for a given platform. Finally, throughput is defined as the total images that can be processed on a given platform in a second, where processed refers to computing the output for an input image.

The experimental results obtained across all three datasets are summarized in Tables 2-4. Considering that the images in the MSTAR and GBSAR datasets are sized $88\times 88$, our experiments involve patch sizes of 8 and 11, along with hidden dimension sizes of 44 and 88, respectively. However, as the image size in the SynthwakeSAR dataset is $128\times 128$, we opt for patch sizes of 8 and 16, accompanied by hidden dimension sizes of 96 and 48, respectively. The depth of the transformer encoder is set to 4, 6, 8, or 12, and the number of heads is configured to be 2 or 4. Each of these hyperparameter settings is employed during the training of the model to assess its performance across different configurations. Latency for FPGA, GPU and CPU platforms is the end-to-end latency defined in 6.5, measured from the time that the input is sent to a processor to the time that the processor takes to compute the output. Throughput for the three platforms is defined as per 6.5 measuring the total images that can be processed in a given second. For latency comparison across the $3$ platforms, we utilize a batch size of $1$. A batch size of $1$ represents a real-time scenario for SAR ATR applications wherein the input images are continuously streaming. In order to compare the throughput for the $3$ platforms, results are computed over batch sizes $1,8,16,32,64,128,256$ and $512$. Note that the FPGA throughput is nominally measured as the inverse of the per-image inference latency, and thus is invariant with respect to batch size.

Through the experiments, we observe that the performance of the VTR model varies across the three datasets for different hyperparameter settings. Notably, the best classification accuracies (highlighted in bold in Tables 2-4) for the MSTAR and SynthWakeSAR datasets (95.96% and 93.47% respectively) are achieved with similar hyperparameter configurations. However, for the GBSAR dataset, a lower depth size results in better performance. It can be concluded that a lower patch size results in a higher classification accuracy as smaller patch sizes are capable of capturing fine-grained details in the input SAR images. It is also worth noting that the SPT and LSA modules are crucial for the performance improvements achieved by the VTR model, as the standard ViT performs poorly on SAR datasets unless pre-trained on significantly larger datasets.

In comparison to the benchmarking conducted by Ye et al.[24] on the three datasets using CNNs, GNNs, and ViTs, our VTR model equipped with the SPT and LSA modules demonstrates improved or similar performance. Particularly, our model outperforms the state-of-the-art for the SynthWakeSAR dataset, achieving an accuracy of 93.47%. For the GBSAR dataset, we achieve a comparable classification accuracy of 99.46%. For the MSTAR dataset, however, GNNs outperform all other models. This is attributed to the fact that images in this dataset contain the actual target ground vehicle object to be classified only within a few pixels in the center of the image. This limits VTRs capability to extract localized features from mostly non-informant tokens, despite the addition of the SPT and LSA modules. In contrast, GNNs excel at capturing this local information, leading to a better performance. These comparison results are summarized in Table 5. We define the best-performing model as the model that achieves the highest classification accuracy. Table 6 compares the number of parameters in the best-performing models. In contrast to previous work, our model has fewer parameters, leading to a smaller model size. Compared against standard pre-trained ViTs with total parameters of the order $10^{8}$, VTR is significantly smaller ($\times 100$) with the largest model size being of the order $10^{6}$ parameters.

The proposed FPGA accelerator, for the best-performing model across the three datasets, has an average speedup (based on single image inference latency) of $\times 30$ and $\times 70$ when compared to GPU and CPU platforms, respectively. For smaller models, the accelerator has a much higher speedup than the nominal speedup (associated with larger models). Thus, it is highly suitable for deployment in real-time SAR ATR workloads. Figure 6 compares the throughput of the best-performing model (for each dataset) versus the baseline FPGA throughput (green line). For smaller batch sizes of $1,8,16$ and $32$, the FPGA throughput is much better compared to its CPU and GPU counterparts. With larger batch sizes, as expected, the GPU throughput overtakes the FPGA throughput. This suggests that despite being optimized for streaming single image input inferences, our proposed FPGA accelerator can also be used for smaller batch sizes such as $8$ or $16$ with a throughput that is still, on average, $\times 2$ that of the GPU throughput.