DRT: A Lightweight Single Image Deraining Recursive Transformer

Since the introduction of the transformer network by Vaswani et al. [19], much research has been done to replace CNNs with transformers as the backbone structure for vision tasks. In particular, Dosovitskiy et al. [6] introduce the Vision Transformer (ViT) by splitting an image into non-overlapping small patches and treating each of the patches as a token in the original transformer. ViT performs well on the classification tasks while maintaining a relatively low computational complexity due to the image patch embedding mechanism. To extend the capability of ViT, Liu et al. [15] propose the Swin-Transformer. One of the purposes behind the design of the Swin transformer is the same as the purpose of our work, to reduce the computational complexity of the vision transformer. This type of transformer performs multi-headed self-attention in local windows instead of globally as in ViT. The computation cost is reduced from $4hwC^{2}+2(hw)^{2}C$ to $4hwC^{2}+2(M)^{2}hwC$, where the image dimension is $C\times h\times w$, and the local window size is denoted by $M$.

Meanwhile, Swin-Transformer employs the shifted window approach to compensate for the loss of global information. A single Swin Transformer Block (STB) firstly passes the input through a layer norm (LN) and then via a windowed multi-headed self-attention (WMSA) mechanism. The output from the self-attention mechanism is combined with the initial input forming the first residual connection. Then the output is passed through a multi-layer perceptron (MLP) followed by another LN. The final output combines again with the output from the self-attention mechanism forming the second residual connection. Multiple STBs are used to process a given image and achieve state-of-the-art results in recognition, detection, and segmentation tasks.

Vision transformers from the last section and their variants have been utilized recently to achieve better general image restoration outcomes. Chen et al.[3] uses a transformer encoder-decoder structure based on [19] to design a general image restoration transformer termed Image Processing Transformer (IPT). Similarly, Wang et al. [20] propose to use a transformer block, namely LeWin block, and an hourglass structure with skip-connections for background reconstructions. This design performs well in different image reconstruction tasks while maintaining a relatively low computational cost. Furthermore, Liang et al. [13] design residual transformer blocks on top of the Swin-Transformer and stack multiple of the residual blocks to perform deep feature extractions for image reconstructions. Unlike IPT, which is pre-trained on ImageNet, this type of transformer is trained with 800 samples while maintaining a small number of network parameters relative to many other CNN-based or transformer-based methods.

A recursive structure in deep learning repetitively calls some portion of layers in a network to process data. The use of the recursive structure for image reconstruction tasks can be traced back to [10]. In this work, a deeply recursive CNN has been proposed to stabilize the training process and maintain a deep structure with few parameters. It empirically shows that a single recursive layer is enough to achieve good image super-resolution results. Later, the same idea is used by [7] to design a residual guided recursive CNN network for image deraining. Our method is similar to this work because we also thoroughly investigate the use of recursive and residual structures on transformer-based architectures instead of relying on CNNs. Although recursive structures can also lead to potential issues such as exploding gradient or overfitting, carefully designed recursions can lead to optimal network architecture, as discussed in the following sections.

DRT consists of three stages, the patch embedding stage, $f_{1}$, the deep feature extraction stage $f_{2}$, and the image reconstruction stage, $f_{3}$. Its visualisation is given in Figure 2. Essentially, one can view the first stage as input pre-processing stage and the third stage as inference output post-processing stage whereas the second stage is the main inference stage. Given a raw rainy image $R_{\text{input}}$, the equation below hence captures the functionality of DRT,

$$\text{DRT}(R_{\text{input}})=f_{3}\circ f_{2}\circ f_{1}(R_{\text{input}})$$

(1)

The patch embedding stage $f_{1}$ firstly applies a convolutional layer to the rainy image $R_{\text{input}}\in\mathbb{R}^{H\times W\times C}$ and map it to $R_{\text{conv}}\in\mathbb{R}^{H//P\times W//P\times D}$ to divide the image into patches and stack them depth-wisely. Here, $H\times W\times C$ is the original image’s height, width and depth respectively. $P$ denotes the patch size used to divide the input into small non-overlapping windows, and $D$ is the embedded dimension and usually $D\gg C$. Note the $//$ sign denotes the integer division operation. Since transformers typically work with one dimensional tokens, we reshape the convolution layer’s output $R_{\text{conv}}\mapsto R_{\text{embed}}\in\mathbb{R}^{HW//P\times D}$ to complete the first stage of the network. This output can then be fed to transformer blocks to process. We can succinctly write the first stage as

$$f_{1}(R_{input})=\text{PatchEmbed}(\text{Conv}(R_{input})).$$

(2)

The next stage consists of multiples of recursive transformer blocks (RTB) composed together to perform deep feature extractions as described by Equation 3. We use $N$ to denote the total number of RTBs. The details of the RTB are given in the next subsection. Here we note that this stage maintains the input dimension, and there is no parameter sharing between each RTB. The output from this stage is denoted as $R_{\text{deep}}$. At the end of the process, a residual connection is used to add the input to the deep feature extraction stage to its output to restrict the RTBs from focusing on detecting rain streaks.

$$f_{2}(R_{\text{embed}})=\text{RTB}_{1}\circ\dots\circ\text{RTB}_{N}(R_{\text{embed}})+R_{\text{embed}}.$$

(3)

The deep features, $R_{\text{deep}}$, are then processed by the image restoration stage, which reverses the process of the first stage as seen in equation 4. The height and width of the input image is firstly restored by reshaping $R_{\text{deep}}\mapsto R_{\text{unembed}}\in\mathbb{R}^{H//P\times W//P\times D}$ and convolution layers are used to process these features to give the output $R_{\text{out}}\in\mathbb{R}^{H\times W\times C}$. Finally, another residual connection is used between the network’s input to its output so that rain streaks features extracted by the network can be removed from the input.

$$f_{3}(R_{\text{deep}})=\text{Conv}(\text{PatchUnembed}(R_{\text{deep}}))+R_{\text{input}}.$$

(4)

Each RTB uses the recursive and residual connections to stack Swin Transformer Blocks (STBs). Mathematically, the STB can be described as,

$$STB(x)=MLP\circ LN(WMSA\circ LN(x)+x)\\ +(WMSA\circ LN(x)+x)$$

(5)

Here the WMSA can be summarised by the equations [19],

$$\text{MultiHead}(Q,K,V)=\text{Concat}(\text{h}_{1}\dots\text{h}_{n})W^{O}$$

(6)

$$\text{h}_{i}=\text{Attention}(QW_{i}^{Q},KW_{i}^{K},VW_{i}^{V})$$

(7)

$$\text{Attention}(Q,K,V)=\text{Softmax}\left(\frac{QK^{T}}{\sqrt{d_{k}}}\right)V$$

(8)

In the above equations, $W$ refers to weights. The inputs, Q, K, and V, often are called queries, keys, and values. The term $\sqrt{d_{k}}$ refers to the dimension of $K$. If we denote the number of STB in each residual connection as $U$ and the number of recursive calls as $L$, let the input be denoted as $S_{\text{in}}\in\mathbb{R}^{HW//P\times D}$, then each RTB can be described with the following equation.

$$RTB(S_{\text{in}})=\bigoplus_{j=1}^{L}\left(\bigoplus_{i=1}^{U}STB_{i}(S_{\text{in}})+S_{\text{in}}\right).$$

(9)

The $\bigoplus$ sign is used to denote multiple function compositions. The dimension of the input and output is kept the same. The residual connection always starts from the input to the RTB. The first $\bigoplus$ in the above equation controls the recursive calls of RTB, and the second $\bigoplus$ controls the composition of different STBs. Each $STB_{i}$ in a residual connection can have different parameters, but recursive calls always invoke the same STBs as seen in the equation above since the index $j$ is never used. For instance, in Figure 2, there are two STBs ($U=2$) in each residual connection, and such connection is recursively called three times ($L=3$) in a single RTB. At the end of each RTB, convolutional layers are used to process information that the local window-based Swin transformer may ignore. If a single convolutional layer is employed, we do not use any activation after the convolution. For multiple convolutional layers, we use Leaky ReLU as activation functions.

There are many proposed loss functions to facilitate the training of the networks; however, we find it is sufficient to use the mean squared error loss. For an output image, $x$ and its corresponding ground-truth image $\bar{x}$, the error is given by,

$$\text{MSE}(x)=\frac{1}{M}\sum_{i=1}^{M}(\bar{x}_{i}-x_{i})^{2},$$

(10)

where M denotes the total number of entries, $\bar{x}_{i}$ denotes the ground truth pixel value of the i^th entry, and $x_{i}$ denotes the prediction value of the i^th entry. For multiple predictions, we simply average the errors for each pair to get a final error measure.

Experiments are carried out on the following data sets to evaluate the efficacy of the proposed DRT.

The PSNRs and SSIMs of the best performing DRT model evaluated on three data sets and the results of comparisons with the other state of the arts are shown in Table 1. Our results are in bold font for better visualization.

For Rain100L, our performance exceeds all methods listed in Table 1 by at least 0.33db under the PSNR metric. For the same metric, our method’s performances ranked second on the Test100 and Rain100H data sets. Similarly, DRT’s results ranked second on the Rain100L under the SSIM metric. On the other hand, our method is very lightweight with respect to the number of parameters used, as seen in Table 1. Specifically, DRT only uses $\sim$ 1.3 % of the number of parameters compared to HiNet. Note for early methods that did not release the source code; we omit their parameters in the table. In general, there is a trade-off between the number of parameters and the derain performance. Refer to Figure 5 for a visualization of this trade-off. Models with a large number of parameters are ignored in this graph for better visualization. Methods closer to the upper-left corner tend to balance the number of parameters and the derain result and vice versa. DRT balances this trade-off well by being the closest to the upper-left corner, whereas methods such as Jorder-E and HiNet obtain good performance with the cost of large network sizes. Other methods may use fewer parameters but cannot achieve derain results as good as DRT. We further provide a fixed input to HiNet, PreNet, and DRT to calculate their corresponding multiplier-accumulator operations (MACs) (this metric is also used by HiNet) and the amount of memory needed to store the network and perform a forward and backward passing. The results are shown in Table 2. DRT consumes the least amount of MACs while enjoying roughly 4.5 times less memory usage than HiNet.

We select visual outcomes of different methods evaluated on Test100, Rain100L and Rain100H and present them in Figure 3, Figure 4 and Figure 1 respectively. These images are chosen to show the deraining performance of different rain types. In Figure 3, the rain density is high, but each rain streak is very small. DRT is able to remove most of these rain streaks, whereas Hi-Net ignores many of them. Note that there is a tone shift in the background of the input and the background of the ground truth (GT) images. Since the training data do not consider background color changes, our restored image cannot recover the original tone. In Figure 4, other methods seem to interpret some rain streaks as parts of the land or the cloud and hence cannot fully remove these rains, whereas DRT can remove almost all of them. In Figure 1, we can see that artifacts are present in the smoke produced from the train, and artifacts exist in the smoke for images produced by Jorder and Hi-Net. Our model and PreNet give a better reconstruction of these complex soft-edge forms. Lastly, we also present a realistic derain outcome as seen in Figure 6. Compared to Hi-Net, DRT can remove most of these rains, as seen in the zoomed-in pictures. Note for the realistic image that the haze caused by rains remains expected since the training data does not consider a mixture of haze and rain.

We found out that the most critical design choice is reducing the patch size used by STBs. The original Swin transformer block sets the patch size to 4. We experience a significant increase in PSNR when reducing the patch size to one. This shows the importance of considering local pixel information in image deraining. Furthermore, we also found it beneficial to keep the input dimension (i.e., do not perform any up-sampling or down-sampling to the input) throughout the entire network. Since single image deraining is essentially a dense prediction task, these two design choices ensure there is no loss of pixel-level information for the network to infer. Other hyper-parameters of the model are also tuned extensively, and we summarise them here. As mentioned before, we did not find any impact on the use of the shift-window approach from the original Swin transformer; this is likely due to the fact that we have used convolutions at the end of each RTB and local positional encoding for each STB to consider both local and global positional information whereas the shift-window approach is proposed to solve the loss of global information in the original transformer. Similarly, we tried to increase the depth of the transformer’s working dimension from 96 to 180, but there is no improvement in performance. We tried to drop the residual connection between the input and the output, but no improvement was observed. This is likely because residual connections limit the inference stage to detect rain streaks, which occupy a small portion of information in the input image.

We then perform an ablation study on the essential design features of the DRT, namely, the number of RTBs ($N$), the number of STBs, and the number of recursive calls on STBs ($L$). During this study, each model is trained on Rain100L from scratch with a similar training procedure. All other hyper-parameters not included in the ablation study are fixed on the optimal network structure. The evaluation of each model is performed on the Rain100L data set. The number of parameters and the PSNR for each evaluation is listed in Table 3 with the best model displayed in the last row. As seen in the table, the first two rows study the impact of the number of RTBs concerning the best DRT model; however, corresponding low PSNR results show that both increasing and decreasing the number of RTBs dramatically limits the performance of DRT. Row three to five studies the influence of the number of recursive calls of STBs on the performance of DRT. Once again, neither increasing nor decreasing the number of calls can boost the derain performance of DRT. Note that having only one recursive call refers to a variant of DRT without any recursion. However, its performance is the worst among all other variants, which shows the importance of recursions in the design. We then tune both the number of RTBs and the number of recursive calls simultaneously, but the result is not good. Hence, we conclude that having six RTBs and three recursive calls is essential for a promising derain performance. Contrasting to this, the second and third last row of Table 3 indicate, that changes in the number of STBs do not influence the derain result much. It is worth noting that having just one STB lowers the number of parameters from 1.18M to below a million, 0.841M, while still keeping a good derain outcome.

We further test DRT’s performance on the task of single image desnowing to empirically show its capability in solving different image restoration tasks. The best model is trained on the Snow100K data set [14] with the same training procedure as before. The model is then tested on the corresponding three data sets provided by Snow100K; each consists of one snow type. The PSNR results of different methods are shown in Table 4. Our model’s performance ranks either the first or the second among these methods. We also display a sample of desnow visual results with the corresponding input, output, and mask images in Figure 7. Although the snow in the input image is dense, DRT can remove most of them and restore a clean background that is very close to the ground truth. Rain streaks and snowflakes can be very different both physically and visually, which means they require different functions to be removed; this study shows that upon training, DRT is not restricted to single image deraining but can also be well applied to other fields.