Memory-Constrained Semantic Segmentation for Ultra-High Resolution UAV Imagery

Semantic segmentation is modeled as a dense prediction task, many works [14, 15, 16, 17, 18, 19, 20] based on convolutional neural network has achieved great success. In recent years, several methods [21, 22, 23, 24, 25, 26] use Transformer architecture to conduct semantic segmentation tasks. However, most of the work is applied to ultra-high resolution images, which raises the trade-off between performance and GPU memory. To this end, Chen et al.[10] integrate global image and local patch each other in the deep layer to balance performance and GPU memory usage. Limited by the speed of global and local interaction, Wu et al.[27] design a classification network to choose important patches for the feature fusion. To further improve performance, Huynh et al.[11] progressively refine coarse segmentation results via a multi-stage processing pipeline. Li et al.[12] introduce a multi-scale locality-aware contextual correlation and the adaptive feature fusion scheme to strengthen local segmentation. These methods are patch-based approaches, which can save GPU memory but consume time due to multiple local segmentation. In the latest work, Guo et al.[28] leverage the shallow-deep network to directly process the full-scale ultra-high resolution images for accelerating the inference speed. In addition, some work [29, 30] generates high-quality semantic results by refining coarse segmentation maps from a pre-trained model. Comparing with the previous works, we focus on the systems constrained by limited GPU memory, by considering a better trade-off between accuracy, GPU memory, and speed.

In implicit neural representation (INR), the signals of the object and scene are maps from coordinates via a multi-layer perceptron (MLP) applied in modeling 3D reconstruction [31, 32, 33, 34, 35, 36]. For example, Mildenhall et al.[35] present NeRF that learns an implicit representation for a novel scene view using multiple image views. It can accurately capture the intricate details of the shape with a minimal amount of parameters. Later, INR is also used in video representation [37, 38, 39]. Chen et al.[37] propose a novel neural representation for videos that takes the time index as input and directly outputs the corresponding RGB frame. Recently, INR has also made some progress in 2D tasks [13, 40, 30, 41]. Chen et al.[13] represent natural and complex images in a continuous manner, which are trained in the image super-resolution task. Xu et al.[40] perform the spatial encoding in implicit functions and further introduce deep coordinate fusion and residual MLP architecture. Hu et al.[41] propose an alignment function using multi-level feature fusion for semantic segmentation. Among them, [41] is most related to our work. The key difference rests in that we utilize high-resolution spatial masks as guidance to the query module for the interpretable details of the ultra-high resolution images, while [41] only utilizes INR to perform multi-scale feature alignment.

Similar to the human brain, deep neural networks encode, store, and extract information via memory. In the vision tasks [42, 43, 44, 45, 46, 47, 48], it can capture cross-image information to serve the current image. Wang et al.[43] use a cross-batch memory mechanism to record and update the embeddings of past iterations for the collection of sufficient hard negative pairs. Xie et al.[46] relieve the ambiguity of similar objects by memorizing and tracking the regions of target objects. In [47], Kim et al.store the domain-agnostic categorical knowledge in the memory to achieve domain generalization for semantic segmentation. Jin et al.[48] set up a memory module to store the dataset-level distribution information of all classes and perform a coarse-to-fine iterative inference strategy in the memory. In our work, we introduce a memory-based interaction scheme that stores low-resolution semantic information to efficiently enhance the spatial semantics of the high-resolution image. Therefore, it is able to rectify spatial-guided semantic bias in our query module.

In this paper, we introduce a semantic segmentation approach capable of efficiently processing ultra-high resolution images on systems with limited GPU memory.

In response to the GPU memory constraint, we follow the patch-based paradigm (e.g., [12]) in which we partition a large image into local patches, segment them subsequently, and then merge all the local segmentation results together. Formally, given an ultra-high resolution image $\dot{I}$ with width $\dot{W}$ and height $\dot{H}$, our approach partitions the image into $N$ overlapping local patches (denoted as $I$ in the following) along both the row and column axes. Then, our model generates the segmentation results for each local patch, which are merged to produce an ultra-high resolution segmentation mask.

Given the possibly large resolution of the image patch $I$ (e.g. 1280 $\times$ 1280), there still arises a necessity to mitigate GPU memory expenses. This, in turn, imposes constraints on the depth of the encoding network, which consequently limits its capability to extract an ample amount of semantic information. Therefore, we downscale $I$ to a properly small resolution (e.g. 320 $\times$ 320) for extracting semantics. To this end, as illustrated in Fig. 1, our proposed framework is based on a three-branch architecture, comprising of the semantic branch, memory-based interaction branch, spatial branch, and spatial-guided high-resolution query module. Specifically, a partitioned image patch $I$ and its downscaled version $\hat{I}$ is fed into the encoders for feature extraction. The spatial branch and semantic branch obtain the visual features and semantic features, respectively. In addition, the memory-based interaction branch relies on an external memory bank to mitigate the bias in the guidance information, which associates the cross-image compact semantic representation with high-resolution spatial information for regularization. Inspired by [13], the spatial-guided high-resolution query module is dedicated to infer the high-resolution semantics in a pixel-by-pixel manner by querying the corresponding latent embeddings with the guidance of high-resolution structural information and low-resolution semantics. In the following, we will elaborate on the technical details of the two components.

High-quality local image segmentation relies on contextual cues beyond local patches. However, incorporating such context can increase computational overhead. To balance segmentation quality and computational efficiency, we propose a novel spatial-guided high-resolution query module that queries the nearest latent embedding of a given spatial coordinate to obtain the corresponding semantic result without the need for additional context.

In concrete, we first define a query function $f_{\theta}$ ($\theta$ is learnable parameters) over the feature maps to achieve the high-resolution semantic mask $S$ ($S\in\mathbb{R}^{C\times H\times W}$) where $C$ denotes the number of semantic classes. Here, the feature maps are viewed as latent embeddings evenly distributed in spatial dimensions and we assign corresponding coordinates to them. Hence, the value at the coordinate $x_{q}$ of the high-resolution semantic mask $S$ can be queried as below:

where $z^{*}$ is the nearest latent embedding from $x_{q}$ and $x^{*}$ is the low-resolution coordinate of the latent embedding $z^{*}$ in the spatial domain. Given the relative coordinates, the high-resolution query function $f_{\theta}$ can query the nearest semantic result set of the latent embedding $z^{*}$.

To this end, we adopt a vanilla MLP as the query function shared by each latent embedding. However, the previous study [40] implies that neural networks are insensitive to high-frequency signals and are inclined to learn low-frequency signals. This may lead to undesirable artifacts for intricate ultra-high resolution images. Consequently, we encode the relative coordinates via a periodic function $\phi$ to enhance the capability of the network in the high-frequency domain. Thus, the encoding function of the relative coordinates (i.e., $x_{q}-x^{*}$) is defined as:

where $\omega_{1},...,\omega_{n}$ are initially set to $2e^{i}$ ($i\in[1,...,n]$) as the frequency parameters that will be fine-tuned on the training stage. As such, $\phi(\cdot)$ maps the relative coordinate to a $2n$-dimensional positional information. Eq. 1 can be rewritten as:

In general, the spatial resolution of the latent embedding is much smaller than that of a semantic map, which results in the loss of spatial details during the query. To this end, we propose to utilize the higher resolution masks to provide spatial guidance for the latent embeddings. Specifically, the high-resolution segmentation masks $M_{b}$ ($M_{b}\in\mathbb{R}^{C\times\frac{H}{2}\times\frac{W}{2}}$) and $M_{l}$ ($M_{l}\in\mathbb{R}^{C\times H\times W}$) are computed by the memory-based interaction branch and spatial branch, respectively, which modifies the high-resolution query function as below:

where $m_{b}^{*}$ and $m_{l}^{*}$ denote the nearest mask values from $x_{q}$ in mask $M_{b}$ and $M_{l}$. $x_{b}^{*}$ and $x_{l}^{*}$ are the corresponding low-resolution coordinates of mask values $m_{b}^{*}$ and $m_{l}^{*}$, separately. With the cues of high-resolution spatial information, the latent embeddings can better predict fine semantic results.

In our model, the high-resolution spatial mask $M_{b}$ is computed by the last layer of the memory-based interaction branch. Since this branch is relatively shallow, its estimated high-resolution mask contains semantic bias that may interfere with spatial cues for the query function. To mitigate the negative impact, it requires contextual information for regularization. Without the need for additional large computation, we propose an efficient memory-based interaction scheme that adds semantics to the mask $M_{b}$. It involves an external memory bank $\mathcal{M}$ ($\mathcal{M}\in\mathbb{R}^{D\times C}$) that stores the semantic features across images, where $D$ denotes the feature dimension. To reduce the computation overhead to the largest extent, our proposed memory scheme costs linear complexity only.

As the preliminary step, we randomly select an image sample from the training set. Then, we calculate a mean vector of the semantic branch features for each category as the initial value of the memory bank, with the aid of the ground-truth segmentation maps. During training, the representation of each category $c$ ($c\in[1,...,C]$) in the memory bank $\mathcal{M}$ is updated by the moving average method in the $t$-th iteration.

where the momentum $m$ is set as 0.9, and $\varphi$ is a transform function. $\mathcal{R}_{t}$ ($\mathcal{R}_{t}\in\mathbb{R}^{D\times\frac{H}{4}\times\frac{W}{4}}$) is the semantic branch features of the current sample in the $t$-th iteration. In $\varphi$, $\mathcal{R}_{t-1}$ is permuted with the dimension $D\times N$ ($N=\frac{HW}{16}$).

The feature representation of each category in the memory bank can be denoted as $\mathcal{R}^{c}$ ($\mathcal{R}^{c}\in\mathbb{R}^{D\times N^{c}}$) that stores the representation of the category $c$, where $N^{c}$ is the number of the pixels labeled as the category $c$. $\mathcal{GT}$ ($\mathcal{GT}\in\mathbb{R}^{D\times N}$) stores the ground-truth category labels of $\mathcal{R}$. Next, we calculate the cosine similarity $\mathcal{S}^{c}$ ($\mathcal{S}^{c}\in\mathbb{R}^{N^{c}}$) between $\mathcal{R}^{c}$ and $\mathcal{M}^{c}$:

Finally, the transform function $\varphi$ outputs:

After updating the memory bank, we associate the memory bank as the compact cross-image semantic representation to high-resolution information to enhance the mask $M_{b}$ in the semantic perspective. Specifically, we read the memory bank $\mathcal{M}$ and extract the features $\mathcal{F}_{b}$ ($\mathcal{F}_{b}\in\mathbb{R}^{D\times\frac{H}{2}\times\frac{W}{2}}$) from the memory-based interaction branch. Then, $\mathcal{F}_{b}$ is permuted into the features with the dimension $D\times\frac{HW}{4}$. Thus, we calculate the relation $\mathcal{W}$:

where $\otimes$ is matrix multiplication. The size of $\mathcal{W}$ is $C\times\frac{HW}{4}$ and it is reshaped as $C\times\frac{H}{2}\times\frac{W}{2}$. Last, $M_{b}$ is refined as follows:

Note that, compared with the previous memory scheme with quadratic complexity (e.g. [48]), our proposed memory scheme has a linear computational complexity of $\mathcal{O}(\frac{HWC}{4})$. This is because our scheme does not associate the global pixel memory for each pixel, which spares the self-attention computation, contributing significantly to a substantial increase in overall computational efficiency.

Inria Aerial [49]. This dataset contains large resolution aerial images of five cities, ranging from dense metropolitan districts to alpine resorts. It contains 180 aerial images of 5000$\times$5000 pixels with the binary mask for building/non-building areas. Following the protocol of [10], we split images into training, validation, and testing sets with 126, 27, and 27 images, respectively.

DeepGlobe [50]. This dataset provides 803 ultra-high resolution aerial images with 2448$\times$2448 pixels. It contains seven classes of landscape regions, including urban, agriculture, rangeland, forest, water, barren, and unknown region not considered in the challenge. We split all images following [10], i.e., 454 training images, 207 validation images, and 142 testing images.

Training details. We implement our framework using the mmSegmentation [51] toolbox on a workstation with a single NVIDIA RTX 3090 GPU. In particular, we adopt DeepLabv3 [52] with ResNet18 [16] as the encoder of the semantic branch and STDC [53] as the encoder of the memory-based interaction branch. During training, we optimize the parameters adapting Stochastic Gradient Descent (SGD) and set the batch size to 4 and 8 for Inria Aerial and DeepGlobe, respectively. The initial learning rate is set to $10^{-2}$, which is decayed by a poly learning rate policy with the power of 0.9. In practice, it takes 40k and 80k iterations to converge our model for two datasets, respectively.

Inference details. During inference, we measure the GPU Memory and Frames-per-second (FPS) on an RTX 2080Ti GPU and adopt the same environment as [28] (i.e., CUDA 10.1, CuDNN 7.6.5, and Pytorch 1.6.0) for fair comparison.

In practice, segmenting an ultra-high resolution image often consumes an exceedingly large amount of GPU memory resources. As a consequence, for the robotic platforms like UAVs, they usually have limited computation resources and GPU memory for segmenting and analyzing the large images on board. To simulate the challenging situations, in the experiments, we compare the segmentation methods under small and large GPU memory limits. In this experiment, we set 1.5 GB as the small GPU memory limit. For large GPU memory limits, we set 7.5 GB for Inria Aerial and 3.5 GB for DeepGlobe, due to their different image resolutions. In the following, we conduct the comparison experiments and discuss the results.

Small memory limits. In general, there are two ways to segment the ultra-high resolution images: 1) segmenting downsampled global images (denoted as global inference); and 2) cropping, segmenting, and merging local patches (denoted as local inference). We first consider the methods that adopt the global inference strategy, which straightforwardly downsamples the whole input image $\dot{I}$ to meet the requirements of small limited GPU memory. To this end, we retrain and test U-Net [54], FCN-8s [15], and DeepLab v3+ [19] on Inria Aerial and DeepGlobe datasets. As shown in Table I and Table II, these methods can hardly achieve satisfactory accuracy despite their high FPS, since compressing large images causes severe detail lost.

Thus, we adopt the local inference strategy for the segmentation methods to meet small GPU memory limits. Table. I and Table. II show our superior performance than other methods in the respective datasets. In specific, our model shows similar GPU memory usage and running time comparing to the state-of-the-art ISDNet [28], but elevates 2.50% mIoU on Inria Aerial and 0.87% mIoU on DeepGlobe. It is worth noting that our model achieves a significant advantage on Inria Aerial dataset, because our proposed spatial-guided high-resolution query module can depict delicate objects and this dataset contains a large number of small buildings. Besides, as a ResNet-50 based model, PPN [27] is also able to fulfill the requirement of limiting small GPU memory, but it underperforms our method and ISDNet. In Fig. 4, we illustrate several qualitative comparison results. We observe that our model has the ability to delineate fine regions (e.g., small urban and rivers) and stronger semantic discrimination (e.g., large forests and rangeland). This is attributed to the high-resolution query and memory-based semantic enhancement.

Large memory limits. To evaluate the capability of our model under large GPU memory limits, we compare our model with BiSeNetV1 [55], STDC [53], PointRend [56], CascadePSP [29], GLNet [10], MagNet [11], FCtL [12], and ISDNet [28], in terms of mIoU, GPU Memory, and FPS. The comparison results are depicted in Table. III and Table. IV, in which we follow most of the results reported by [28]. As observed, our model is comparable to the state-of-the-art ISDNet on DeepGlobe, but we gain at least 3% improvement than it on Inria Aerial. This shows that our approach can further trade off performance, GPU memory, and speed.

Performance-memory trade-off. We quantify and visualize all the state-of-the-art methods based on the small and large GPU memory limits in Fig. 2. We can see that our model can achieve optimal mIoU regardless of small or large GPU memory limits. More importantly, our model can obtain similar performances under small and large GPU memory limits (77.25% v.s. 77.27% on Inria Aerial and 73.66% v.s. 73.50% on DeepGlobe ), indicating that our robustness to GPU memory constraints, whereas other methods are greatly affected due to narrow view of local patches. This still benefits from our query module, which predicts pixel categories only by the corresponding relative coordinates and the nearest latent embeddings.

In the following, we conduct ablation studies on the two proposed modules. First, we demonstrate the effectiveness of spatial-guided high-resolution information and memory-based interaction scheme in our model. Next, we analyze different strategies for updating and reading the memory bank in our scheme. Last, we investigate the boundary case of GPU memory usage.

Model structure. First, we take the naive query module as our baseline, as shown in Table V, our model just costs 1284MB GPU memory with high FPS, but only gains 69.24% mIoU. To improve performance, we use STDC [53] to extract high-resolution spatial information $M_{b}$. It is able to boost the performance from 69.24% to 72.06%, which exceeds bilinear interpolation with 71.81% mIoU and demonstrates the effectiveness of the spatial mask guidance. Besides, the GPU memory increase is only 186MB and the total GPU memory is still less than 1.5GB. Based on this, we add memory bank $\mathcal{M}$ and higher resolution information $M_{l}$, separately. It is observed that there is a great performance improvement, especially the guide of spatial information $M_{l}$ (at least 1.1% boost). Importantly, the GPU memory and time hardly grow (about 1MB and less than 0.5 FPS) because both are linear operations. Next, with the aid of both, our model can achieve the best result (i.e., 73.66% mIoU and 1472MB GPU Memory). Besides, we also try to enhance the semantics of higher resolution masks $M_{l}$ using a new memory bank $\mathcal{M}_{l}$, but the result gets worse, we think that this mask mainly contains spatial details and the semantics are too weak to improve. Particularly, we can also adopt high-resolution features as the spatial information guide, that is implicit feature alignment (IFA) [41]. Table V shows IFA is inferior to our method by at least 1.4% mIoU and uses 270MB more GPU memory than ours. This largely reflects that our method is more effective and efficient.

Memory-based interaction scheme. In our memory scheme, we use semantic branch features to update the semantic memory, dubbed cross-branch interaction strategy. To demonstrate the necessity of this strategy, we perform ablation experiments with regard to it. Table VI shows performance barely changes (73.23% v.s. 73.25%) using the features of the memory-based branch, compared with no memory bank. For update mode, we ablate the effect of ”Mean” mode (i.e., $\hat{\mathcal{R}}^{c}=\sum^{N^{c}}_{i=1}\frac{\mathcal{R}^{c}_{i}}{N^{c}}$). As shown in Table VI, this is only 0.06% mIoU improvement over no memory bank. Besides, we replace Concat instead of Softmax (i.e., $M_{b}=\text{Concat}(\frac{\mathcal{M}^{T}\otimes\mathcal{F}_{b}}{\sqrt{D}},M_{b})$) for the reading mode, it leads to 0.08% mIoU degradation.

Memory usage boundary. To investigate the GPU memory usage boundary, we attempt to reduce the size of local patches in order to further decrease GPU memory consumption during the inference process. We perform this experiment on Inria Aerial with 5000×5000 pixel images. As shown in Fig. 3, our model and ISDNet achieve similar minimum GPU memory usage of approximately 940MB. However, even under the lowest-bound of GPU memory constraint, our model outperforms ISDNet by a significant margin, with comparatively less degradation in performance. In general, our model demonstrates superior performance over ISDNet, both under small and large GPU memory constraints.