SideRT: A Real-time Pure Transformer Architecture for Single Image Depth Estimation

Our proposed SideRT has a simple yet efficient encoder-decoder architecture that predicts depth from a single image. We adopt Swin Transformers Liu et al. (2021) as our backbones, each image is divided into several $4\times 4$ non-overlapping patches. The patch, along with relative positional embedding, is embedded by the linear projection layer and the result will be processed by a series of shift window based transformer blocks. Feature maps from all four stages will be utilized during decoding.

The input to the decoder is a set of multi-scale feature maps of four stages: (I) stage 1, $\frac{H}{4}\times\frac{W}{4}\times C$, (II) stage 2, $\frac{H}{8}\times\frac{W}{8}\times 2C$, (III) stage 3, $\frac{H}{16}\times\frac{W}{16}\times 4C$, and (IV) stage 4, $\frac{H}{32}\times\frac{W}{32}\times 8C$. $H$ and $W$ respectively represent the height and width of input images, $C$ denotes the number of channels of features from stage 1. Analogous with a standard feature pyramid, the decoder fuses feature maps progressively in a coarse-to-fine way. Our decoder consists two basic modules: cross-scale attention modules and multi-scale refinement modules (i.e. CSA and MSR modules in Figure 2). To obtain global context, CSA modules aim to fuse feature maps following the guidance of semantic similarity, while MSR modules aim to fuse feature maps according to spatial corresponding relationship. Fusion operations are conducted in a coarse-to-fine manner to get the final prediction, which keeps the same resolution as the input images.

Our proposed CSA module consists of two parts: a linear layer to project each feature to the same number of channels and an attention-based fusion to fuse feature maps from adjacent scales according to semantic similarity. Denote input feature maps as $F_{1}\in C_{1}\times H_{1}\times W_{1}$ and $F_{2}\in C_{2}\times H_{2}\times W_{2}$, while $F_{1}$ is from the shallower stage and $F_{2}$ is from the deeper stage of the encoder. After linear projection, we propose an attention-based method to fuse these two feature maps:

where $L(\cdot)$ is the linear projection operation. The linear layer will project $F_{1}$ from $C_{1}\times H_{1}\times W_{1}$ to $C\times(H_{1}\!\times\!W_{1})$, and $F_{2}$ from $C_{2}\times H_{2}\times W_{2}$ to $C\times(H_{2}\!\times\!W_{2})$ respectively.

For traditional CNN-based methods, global context information only exists near the encoder bottleneck and will be gradually weakened during the hierarchical upsampling of the decoder. Different from them, this CSA module is able to naturally capture global context dependency between feature maps from adjacent scales through calculating attention scores $softmax(L(F_{1})\times L(F_{2}))$ across the entire feature map. Furthermore, this global context dependency contains both semantic-related and depth-related information, as illustrated in the visual analysis of Section 5.

Several multi-scale refinement (MSR) modules work collaboratively in a top-down hierarchy pyramid pattern. A coarser-resolution feature map from the higher MSR pyramid level and a finer-resolution feature map from the lower CSA pyramid level are fed into a MSR module. The MSR module outputs a refined feature map with the help of a low-level semantic but more accurately-localized feature map from the CSA pyramid. The output feature map of the MSR module and the next finer-resolution feature map from the lower CSA pyramid level will be fed into the next MSR module until the final feature map is generated. In a word, the whole refinement process is iterated in a coarse-to-fine way.

Each multi-scale refinement module includes three parts: a bilinear interpolation to upsample coarser-resolution feature map from MSR pyramid, an element-wise addition between the upsampled feature map and finer-resolution feature map from the CSA pyramid, and two MLP layers to reduce aliasing effect of upsampling and generate a depth map at each scale. The last two MSR modules are appended to generate the final result with the input resolution. In particular, unlike other MSR modules, they do not receive outputs from CSA modules.

Different from the fusion style of CSA modules which fuse features according to semantic similarity, MSR modules leverage upsampling to fuse features with corresponding spatial positions. Furthermore, the introduction of MLP layers facilitates information to flow globally.

In order to ease the training of the early stages of the whole architecture, we propose a multi-stage supervision (MSS) strategy at each scale to supervise the training process. In each stage, the loss between the prediction and the corresponding ground truth is calculated. Finally, multi-stage losses are weighted summed together.

Due to the limitation of 3D sensors, the depth data is dense at close areas whereas very sparse in the distance. To alleviate this imbalance problem, we adopt the square root loss function introduced in Song et al. (2021). This loss calculates the difference of predicted depth value and the ground truth in the log space, as shown below:

where $y$ and $y^{*}$ respectively represent the predicted depth map and the ground truth, $V$ is the set of valid pixels in the depth map, $N_{V}$ indicates the total number of valid pixels, and the balance coefficient lambda is set to 0.85.

The model is implemented with PyTorch. Swin transformers pretrained on ImageNet Deng et al. (2009) are used as encoders. The proposed model is trained from scratch for 160 epochs with a batch size of 6 through the AdamW optimizer. The weight decaying factor is set to $5e^{-4}$ and the learning rate is set to $1e^{-4}$. It takes 8 NVIDIA A100 SXM GPUs in the training process. Data augmentation is performed during the training phase to avoid overfitting problems. Images from KITTI and NYU are randomly cropped to 704×352 pixels and 512×416 pixels respectively. In addition, we randomly adjust the scale factor, rotate the input color images within a specific range and flip input images horizontally with a probability of 0.5. The speed of all the methods is tested on a single NVIDIA GeForce RTX 2080 Ti GPU.

Two popular datasets (KITTI and NYU) are used for performance evaluation. The KITTI Geiger et al. (2013) dataset contains the road environment acquired in the autonomous driving scene. The resolution of the acquired images is 1242×375 pixels. We adopt the split strategy introduced by Eigen et al. (2014) for performance comparison. The test set contains 697 images from 29 scenes, and the training set contains 23488 images from 32 scenes. The maximum value of prediction depth is 80 meters.

The NYU dataset Silberman et al. (2012) contains 120K images obtained by Microsoft Kinect camera, including 464 indoor scenes with a resolution of 640×480. We also follow Eigen et al. (2014) to set train/test split, which includes 249 scenarios for training and 654 images from remaining 215 scenarios for testing. The depth map is center-cropped into 561×427 when evaluating the performance.

In order to quantify the performance of our model, we used metrics provided by Eigen et al. (2014), which are widely applied in the performance evaluation of monocular depth estimation. The specific formulations of these metrics are shown in table 3.

We compared our method with state-of-the-art methods on the KITTI and NYU datasets on these metrics. As shown in Table 1 and 2, it is obvious that our method achieves the best performance on all metrics. On the KITTI dataset, compared with the previous state-of-the-art, AbsRel has decreased by 6.9% and SqRel has decreased by 8.9%. On the NYU dataset, compared with the previous state-of-the-art, AbsRel decreased by 9.7% and RMSE decreased by 8.0%.

Our small models like SideRT(Swin-T) and SideRT(Swin-S) outperform many state-of-the-art methods in these two challenging benchmarks. Moreover, we find that the usage of heavy backbone like Swin-L with about 200 million parameters do not slow down the inference process and shows similar inference speed with smaller models like SideRT(Swin-B) and SideRT(Swin-S).

It is worth noting that while surpassing state-of-the-art, we also achieved real-time prediction with a speed of 51.3 FPS. For a better understanding of our performance, we visualize predicted depth maps respectively in KITTI and NYU datasets, as shown in Figure 3 and 4. It can be seen that our method can also successfully predict finely detailed object boundaries that other methods cannot predict clearly.

To get a better understanding of the contribution of proposed components to the overall performance, an ablation study is performed in Table 4. All the experiments are done on the KITTI dataset and use Swin-T as backbone. The training and testing strategies are kept the same with Section 4.3.

Cross-scale Attention. To evaluate the importance of CSA modules, we remove them from the overall architecture (Row 4 and 5 in Table 4). It can be observed that the CSA module improves the performance with a small computational overhead (+0.8M Params) and a small decrease in speed (-12.0 FPS). Row 1-2 in Table 4 show that the CSA module significantly improves the depth prediction directly from the encoder Swin-T. In further visual analysis, we find that the receptive field of the encoder is relatively small, and proposed CSA module will enlarge it tremendously. It is worth mentioning that comparing performances of Row 2 in Table 4 and Row 1-2 in Table 1, the simplest architecture (Swin-T+CSA) gets similar performance with VNL Yin et al. (2019) and outperforms some metrics of DORN Fu et al. (2018).

Multi-stage Supervision. As mentioned in previous works Tolstikhin et al. (2021); Dosovitskiy et al. (2020); Liu et al. (2021), the training of a pure transformer architecture is much more difficult than a CNN-based counterpart. After using multi-stage supervision (MSS), we observe that the training loss curve becomes smoother and the model converges more easily. MSS brings a substantial improvement, including all the metrics, showing that this scheme is very appropriate for accelerating training and promoting the performance of pure transformer architectures.

Multi-scale Refinement. As shown in Row 3-5 in Table 4, introducing MSR modules clearly gives a boost to the performance. Theoretically speaking, CSA and MSR modules augment original feature maps from encoders in a collaborative way. CSA focuses on merging features with high similarity from a global prospect. MSR aims to fuse features with similar positions at different pyramid levels.

Inference Speed. Table 4 shows that most of the parameters come from the backbone since our lightweight decoder only contains 1.1 million parameters. After adding our proposed decoder, the AbsRel metric decreases by 74.9%, with the inference speed only decreasing by 20.4%.

Visual Analysis. In order to see whether our proposed CSA module actually enlarges the receptive field of the backbone, we follow the common practices Yang et al. (2018); Fu et al. (2019) to visualize the empirical receptive field size of a CSA module, as shown in Figure 5. For an input image, we select a reference pixel (denoted by red dot) and compute its feature similarities with all the other positions.

It is obvious to see that after adding CSA, reference pixels get a stronger response from a larger scope. We find that CSA modules lead to higher responses from pixels which not only belong to similar categories but also share similar depths. This is a very beneficial attribute for depth prediction tasks.