Collaborative Global-Local Networks for Memory-Efficient Segmentation of Ultra-High Resolution Images

Introduction

With the advancement of photography and sensor technologies, the accessibility to ultra-high resolution images has opened new horizons to the computer vision community and increased demands for effective analyses. Currently, an image with at least 2048$\times$1080 ($\sim$2.2M) pixels are regarded as 2K high resolution media . An images with at least 3840$\times$1080 ($\sim$4.1M) pixels reaches the bare minimum bar of 4K resolution , and 4K ultra-high definition media usually refers to a minimum resolution of 3840$\times$2160 ($\sim$8.3M) . Such images come from a wide range of scientific imaging applications, such as geospatial and histopathological images. Semantic segmentation allows better understanding and automatic annotations for these images. During the segmentation process, the image is pixel-wise parsed into different semantic categories, such as urban/forest/water areas in a satellite image, or lesion regions in a dermoscopic image. Segmentation of ultra-high resolution images plays important roles in a wide range of fields, such as urban planning and sensing , as well as disease monitoring .

The recent development of deep convolutional neural networks (CNNs) has made remarkable progress in semantic segmentation. However, most models work on full resolution images and perform dense prediction, which requires more GPU memories comparing to image classification and object detection. This hurdle becomes significant when the image resolution grows to be ultra high, leading to the pressing dilemma between memory efficiency (even feasibility) and segmentation quality. Table 1 lists a handful of existing ultra-high resolution segmentation datasets: DeepGlobe , ISIC , and Inria Aerial , in comparison to a few classical normal resolution segmentation datasets, to illustrate their drastic differences that result in new challenges. A more detailed discussion of the three ultra-high resolution datasets will be presented in Section 2.3.

Among the extensive research efforts on semantic segmentation, only limited attention have been devoted towards ultra-high resolution images. Typical ad-hoc strategies, such as downsampling or patch cropping, will result in the loss of either high-resolution details or spatial contextual information (see Section 3.1 for visual examples). Our in-depth studies show that high-accuracy methods like FCN-8s and SegNet requires 5GB to 10GB of GPU memory to segment one 6M-pixel ultra-high resolution image during inference. These methods fall into the top-right area in Fig. LABEL:deep_globe_acc_mem(a) with high accuracy and high GPU memory usage. Contrarily, recent fast segmentation methods like ICNet , whose memory usage is much alleviated, drops in its accuracy. These methods locate in the lower left corner in Fig. LABEL:deep_globe_acc_mem(a). Further studies with different sizes of global images and local patches (Fig. LABEL:deep_globe_acc_mem (b) and (c)) prove that typical models fail to achieve a good trade-off between the accuracy and the GPU memory usage.

This paper tackles memory-efficient segmentation of ultra-high resolution images, which presents the first dedicated analysis of this new topic to our best knowledge. The performance aim will be not only segmentation accuracy, but also reduced memory usage, and eventually, the trade-off between the two.

Our proposed model, named Collaborative Global-Local Networks (GLNet), integrates both global images and local patches, for both training and inference. GLNet has a global branch and a local branch, handling downsampled global images and cropped local patches respectively. They further interact and “modulate” each other, through deeply shared and/or mutually regularized features maps across layers. This special design enables our GLNet the capability of well-balancing its accuracy and GPU memory usage (red dots in Fig. LABEL:deep_globe_acc_mem). To further resolve the class imbalance problem that often occurs, e.g., when one is primarily interested in segmenting small foreground regions, we provide a coarse-to-fine variant of our GLNet, where the global branch provides an additional bounding box localization. The GLNet design enables the seamless integration between global contextual information and necessary local fine details, balanced by learning, to ensure accurate segmentation. It meanwhile greatly trims down the GPU memory usage, as we only operate on downsampled global images plus cropped local patches; the original ultra-high resolution image is never loaded into the GPU memory. We summarize our main contributions as follows:

We develop a memory-efficient GLNet for the emerging new problem of ultra-high resolution image segmentation. The training requires only one 1080Ti GPU and inference requires less than 2GB GPU memory, for ultra-high resolution images of up to 30M pixels.

GLNet can effectively and efficiently integrate global context and local high-resolution fine structures, yielding high-quality segmentation. Either local or global information is proven to be indispensable.

We further propose a coarse-to-fine variant of GLNet to resolve the class imbalance problem in ultra-high resolution image segmentation, boosting the performance further while keeping the computation cost low.

Related Work

Fully convolutional network (FCN) was the first CNN architecture adopted for high-quality segmentation. U-Net used skip-connections to concatenate low-level feature to high-level ones, with an encoder-decoder architecture. Similar structures were also adopted by DeconvNet and SegNet . DeepLab used dilated convolution to enlarge the field of view of filters. Conditional random fields (CRF) were also utilized to model the spatial relationship. Unfortunately, these models will suffer from prohibitively high GPU memory requirements when applied to ultra-high resolution images (Fig. LABEL:deep_globe_acc_mem).

As semantic segmentation grows important in many real-time/low-latency applications (e.g. autonomous driving), efficient or fast segmentation models have recently gained more attention. ENet used an asymmetric encoder-decoder structure with early downsampling, to reduce the floating point operations. ICNet cascaded feature maps from multi-resolution branches under proper label guidance, together with model compression. However, these models were not customized for nor evaluated on ultra-high resolution images, and our experiments show that they did not achieve sufficiently satisfactory trade-off in such cases.

2 Multi-Scale and Context Aggregation

Multi-scale has proven to be powerful for segmentation, via integrating high-level and low-level features to capture patterns of different granularity. In RefineNet , a multi-path refinement block was utilized to combine multi-scale features via upsampling lower-resolution features. adopted a Laplacian pyramid to utilize higher-level features to refine boundaries reconstructed from lower-resolution maps. Feature Pyramid Networks (FPN) progressively upsampled feature maps of different scales and aggregated them in a top-down fashion. Hierarchical Auto-Zoom Net (HAZN) utilized a two-step automatic zoom-in strategy to pass the coarse-stage bounding box and prediction scores to the finer stage.

Context aggregation also plays a key role in encoding the local spatial neighborhood, or even non-local information. Global pooling was adopted in ParseNet to aggregate different levels of context for scene parsing. The dilated convolution and ASPP (atrous spatial pyramid pooling) module in DeepLab helped enlarge the receptive field without losing feature map resolution too fast, leading to the aggregation of global contexts into local information. Similar goal was accomplished by the pyramid pooling in PSPNet . In ContextNet , BiSeNet and GUN , the deep/shallow branches were combined to aggregate global context and high-resolution details. considered the contextual information as a long-range dependency modeled by RNNs. It is worth noting that in our GLNet, the context aggregation is adopted in both input level (global/local branch) and feature level.

3 Ultra-high Resolution Segmentation Datasets

We summarize three public datasets with ultra-high images (studied in Section 4). Basic information and visual examples are shown in Table 1 and Fig. 1, respectively.

The DeepGlobe Land Cover Classification dataset (DeepGlobe) is the first public benchmark offering high-resolution sub-meter satellite imagery focusing on rural areas. DeepGlobe provides ground truth pixel-wise masks of seven classes: urban, agriculture, rangeland, forest, water, barren, and unknown. It contains 1146 annotated satellite images, all of size 2448$\times$2448 pixels. DeepGlobe is of significantly higher resolution and more challenging than previous land cover classification datasets.

The International Skin Imaging Collaboration (ISIC) dataset collects a large number of dermoscopy images. Its subset, the ISIC Lesion Boundary Segmentation dataset, consists of 2594 images from patient samples presented for skin cancer screening. All images are annotated with ground truth binary masks, indicating the locations of the primary skin lesion. Over 64% images have ultra-high resolutions: the largest image has 6682$\times$4401 pixels.

The Inria Aerial Dataset covers diverse urban landscapes, ranging from dense metropolitan districts to alpine resorts. It provides 180 images (from five cities) of 5000$\times$5000 pixels, each annotated with a binary mask for building/non-building areas. Different from DeepGlobe, it splits the training/test sets by city instead of random tiles.

Collaborative Global-Local Networks

For training and inference on ultra-high resolution images with limited GPU memory, two ad-hoc ideas may come up first: downsampling the global image, or cropping it into patches. However, they both often lead to undesired artifacts and poor performance. Fig. 2(1) displays a 2448$\times$2448-pixel image, with its ground-truth segmentation in Fig. 2(2): yellow represents “agriculture”, blue “water”, and white “barren”. We then trained two FPN models: one with all images downsampled to 500$\times$500 pixels, the other with cropped patches of size 500$\times$500 pixels from original images. Their predictions are displayed in Fig. 2(3) and (4), respectively. One can observe that the former suffers from “jiggling” artifacts and inaccurate boundaries, due to the missing details from downsampling. In comparison, the latter has large areas misclassified. Note that “agriculture” and “barren” regions often visually look similar (zoom-in panels (a) and (b) in Fig. 2(1)). Therefore, the patch-based training lacks spatial contexts and neighborhood dependency information, making it difficult to distinguish between “agriculture” and “barren” using local patches only. Finally, we provide our GLNet’ prediction in Fig. 2(5) for reference: it clearly shows the advantages of leverage merits from both global and local processing.

2 GLNet Architecture

During the segmentation process, the feature maps from all layers of either branch are deeply shared with the others (Section 3.2.2). Two sets of high-level feature maps are then aggregated to generate the final segmentation mask via a branch aggregation layer $f_{\text{agg}}$ (Section 3.2.3). To constrain the two branches and stabilize training, a weakly-coupled regularization is also applied to the local branch training.

2.2 Deep Feature Map Sharing

To collaborate with the local branch, feature maps from the global branch are first cropped at the same spatial location of the current local patch and then upsampled to match the size of the feature maps from the local branch. Next, they are concatenated as extra channels to the local branch feature maps in the same layer. In a symmetrical fashion, the feature maps from the local branch are also collected. The local feature maps are first downsampled to match the same relative spatial ratio as the patches were cropped from the large source image. Then they are merged together (in the same order as the local patches were cropped) into a complete feature map of the same size as the global branch feature map. Those local feature maps are also concatenated as channels to the global branch feature maps, before feeding into the next layer.

Fig. 4 illustrates the process of deep feature map sharing, which is applied layer-wise except the last layer of branches. The sharing direction can be either unidirectional (e.g. sharing global branch’s feature maps to local branch, $\mathcal{G}\rightarrow\mathcal{L}$) or bidirectional ($\mathcal{G}\rightleftarrows\mathcal{L}$). At each layer, the current global contextual features and local fine structural features take reference and are fused to each other.

2.3 Branch Aggregation with Regularization

The two branches will be aggregated through an aggregation layer $f_{\text{agg}}$, implemented as a convolutional layer of 3$\times$3 filters. It takes the high-level feature maps from the local branch’s $L^{\text{th}}$ layer $\bm{\hat{X}}_{L}^{\text{Loc}}$, and same ones from the global branch $\bm{\hat{X}}_{L}^{\text{Glb}}$, and concatenate them along the channel. The output of $f_{\text{agg}}$ will be the final segmentation output $\bm{\hat{S}}^{\text{Agg}}$. In addition to the main segmentation loss enforced on $\bm{\hat{S}}^{\text{Agg}}$, we also apply two auxiliary losses, to enforce the segmentation output from the local branch $\bm{\hat{S}}^{\text{Loc}}$ and from the global branch $\bm{\hat{S}}^{\text{Glb}}$ to be close to their corresponding segmentation maps (local patch / global downsampled), respectively, which we find helpful for stabilizing the training.

We find in practice that the local branch is prone to overfitting some strong local details, and “overriding” the learning of global branch. Therefore, we try to avoid the local branch from learning “too much faster” than the global one, by adding a weakly-coupled regularization between feature maps from the last layers of two branches. Specifically, we add the Euclidean norm penalty $\lambda\|\bm{\hat{X}}_{L}^{\text{Loc}}-\bm{\hat{X}}_{L}^{\text{Glb}}\|_{2}$ to discourage large relative changes between $\bm{\hat{X}}_{L}^{\text{Loc}}$ and $\bm{\hat{X}}_{L}^{\text{Glb}}$ with $\lambda$ empirically fixed as 0.15 in our work. This regularization is mainly designed to make local branch training “slow down” and more synchronized with global branch learning, and it only updates the parameters in local branch.

3 Coarse-to-Fine GLNet

For segmentation to separate foreground and background (i.e., binary masks), the foreground often takes little space in ultra-high resolution images. Such class imbalance may seriously damage the segmentation performance. Taking the ISIC dataset for example, $\sim$99% of images have more background than foreground pixels, and over 60% of images have less than 20% foreground pixels (see the blue bars in Fig. 7(1)). Many local patches will contain nothing but background pixels, which leads to ill-conditioned gradients.

To alleviate the class imbalance, we propose a novel two-stage coarse-to-fine variant of GLNet (Fig. 5). It first applies the global branch alone to fulfill a coarse segmentation on downsampled images. A bounding box is then created for the segmented foreground regionIn practice, we dynamically relax the bounding box size, so that the bounded region has a foreground-background class ratio around 1, to have class balance for the second step.. The bounded foreground in the original full-resolution image is then fed as the input for the local branch for fine segmentation. Different from GLNet admitting parallel local-global branches, this Coarse-to-Fine GLNet admits a sequential composition of the two branches, where the feature maps only within the bounding box are first deeply shared from the global to the local branch during the bounding box refinement, and then shared back. All regions beyond the bounding box will be predicted as the background. The Coarse-to-Fine GLNet also reduces the computation cost, through selective fine-scale processing.

Experiments

In this section, we evaluate the performance of GLNet on the DeepGlobe and Inria Aerial datasets and evaluate the efficacy of the coarse-to-fine GLNet on the ISIC dataset. We thoroughly compare our models with other methods to show both the segmentation quality and memory efficiencyWe have chosen several state-of-the-art models with public implementations for comparison (see supplementary for more elaborations). Ablation study is also carefully presented.

In our work, we adopt the FPN (Feature Pyramid Network) with ResNet50 as our backbone. The deep feature map sharing strategy is applied on the feature maps from conv2 to conv5 blocks of ResNet50 in the bottom-up stage, and also on the feature maps from the top-down and smoothing stages in the FPN. For the final lateral connection stage in the FPN, we adopted the feature map regularization, and aggregate this stage for the final segmentation. For simplicity, both the downsampled global image and the cropped local patches share the same size, 500$\times$500 pixels. Neighboring patches have 50-pixel overlap to avoid boundary vanishing for all the convolutional layers. We use the Focal Loss with $\gamma=6$ as the optimization target for both the main and two auxiliary losses. Equal weights (1.0) are assigned to the main and auxiliary losses. The feature map regularization coefficient $\lambda$ is set to 0.15.

To measure the GPU memory usage of a model, we use the command line tool “gpustat”, with the minibatch size of 1 and avoid calculating any gradients. Note that only a single GPU card is used for our training and inference.

We conduct experiments using the PyTorch framework . We use the Adam optimizer ($\beta_{1}=0.9$, $\beta_{2}=0.999$) with learning rate of $1\times 10^{-4}$ for training the global branch, and $2\times 10^{-5}$ for the local branch. We use a minibatch size of 6 for all training. All experiments are performed on a workstation with NVIDIA 1080Ti GPU cards.

2 DeepGlobe

We first apply our framework to the DeepGlobe dataset. This dataset contains 803 ultra-high resolution images (2448$\times$2448 pixels). We randomly split images into training, validation and testing sets with 455, 207, and 142 images respectively. The dense annotation contains 7 classes of landscape regions, where one class out of seven called “unknown” region is not considered in the challenge.

To evaluate the performance of our global-local collaboration strategy, we progressively upgrade our model from shallow to deep feature map sharing (Table 2). With downsampled global images or image patches alone, each branch could only achieve the mean intersection over union (mIoU) of 57.3% and 66.4% respectively. By the aggregation of the high-level feature maps from two branches and the regularization between them, the performance can be boosted to 70.3%. When we share only a single layer of feature maps from the global to the local branch (“shallow sharing”), the aggregated results increased by 0.2%, and when we upgrade to “deep sharing” where feature maps of all layers are shared, the mIoU is rocked to 70.9%. Finally, the bidirectional deep feature map sharing between two branches enables the model to yield a high mIoU of 71.6%.

This ablation study proves that, with deep and diverse feature map sharing/regularization/aggregation strategies, the global and local branch can effectively collaborate together. It is worth noting that even with the bidirectional deep feature map sharing approach (last row in Table 2), the memory usage during inference is only slightly increased from 1189MB to 1865MB.

Fig. 6 visualizes the achieved improvements with two zoom-in panels (a) and (b) showing details. There are undesired grid-like artifacts and inaccurate boundaries in the global (Fig. 6(3)) or local results (Fig. 6(4)) alone. From aggregation, shallow feature map sharing, and finally to bidirectional deep feature map sharing, progressive improvements can be observed with both significantly reduced misclassification and inaccurate boundaries.

2.2 Accuracy and memory usage comparisonwe use public available segmentation models [42, 43]

Models trained and inferenced with global images or local patches may yield different results. This is because models have different receptive fields, convolution kernel sizes, and padding strategies, which results in different suitable training/inference choices. Therefore we carefully compared models trained with these two approaches in this ablation study. We train and test a model twice (with global images or local patches each time), and then pick its best result.

Table 3 shows that all models achieve higher mIoU under global inference, but consume very high GPU memories. Their memory usages drop in patch-based inference, but accuracies also nose dive. Only our GLNet achieves the best trade-off between mIoU and GPU memory usage. We plot the best achievable mIoU of each method in Fig. LABEL:deep_globe_acc_mem(a).

We select FCN-8s and ICNet for in-depth evaluations with different image/patch sizes, since they achieve high mIoU and efficient memory usage respectively. We plot details of this ablation study in Fig. LABEL:deep_globe_acc_mem(b) and (c). For both FCN-8s and ICNet, higher accuracy means to sacrifice GPU memory usage, and vice versa. This proves that typical models fail to balance their segmentation quality and efficiencyIn training with large global images, the minibatch size is limited by the heavy memory usage. We adopt the “late update” optimization trick, e.g. a minibatch size of 2 with weights updated every three minibatches..

The ISIC Lesion Boundary Segmentation Challenge dataset contains 2594 ultra-high resolution images. We randomly split images into training, validation and testing sets with 2074, 260, and 260 images respectively.

On the heavily imbalanced ISIC dataset, the global and local branch can only achieve 72.7% and 48.5% mIoU respectively. When we apply our coarse-to-fine strategy, we can clearly see a much more balanced foreground-background class ratio (red bars in Fig. 7(1)).

By cropping a relaxed bounding box for the foreground (Section 4.2), the local branch is only trained on smaller and class-balanced images, and the cropped-out margin region is assumed as background by default. With class-balanced images, the global-to-local sharing strategy yields 73.9% mIoU, and further bidirectional sharing boosts the performance to 75.2%. In this scenario, the global branch uses a more accurate global context since there is less information loss during downsampling the cropped smaller images. This success proves that the coarse-to-fine segmentation can better capture the context information and solve the class-imbalance problem. We list the results of this ablation study in Table 4 and some visual results in Fig. 7(2).

We finally list mIoU and inference memory usage of GLNet on the local test set of the ISIC dataset in Table 5. GLNet yields mIoU 75.2% and is quantitatively better than other methods on both accuracy and memory usage.

4 Inria Aerial

The Inria Aerial Challenge dataset contains 180 ultra-high resolution images, each with 5000$\times$5000 pixels. We randomly split images into training, validation and testing sets with 126, 27, and 27 images respectively. Table 6 demonstrates the efficacy and efficiency of our deep feature map sharing strategy. The testing results are listed in Table 7, where our proposed GLNet yields mIoU 71.2%. Again, GLNet is quantitatively better than other methods on both accuracy and memory usage. It is worth noting that our GLNet preserves a low memory usage even for a “super” ultra-high resolution image with 5000$\times$5000 pixels.

Conclusions

We proposed a memory-efficient segmentation model GLNet specifically for the ultra-high resolution images. It leverages both the global context and local fine structure effectively to enhance the segmentation in the scenario of ultra-high resolution without sacrificing the GPU memory usage. We also proved that the class imbalance problem can be solved by our coarse-to-fine segmentation approach.

We believe that pursuing an optimal balance of GPU memory and accuracy is essential for the study of ultra-high resolution images, which makes our model important. Our work is pioneering this new research topic of memory-efficient segmentation of the ultra-high resolution images.

Acknowledgement

The work of Z. Wang is in part supported by the National Science Foundation Award RI-1755701. The work of X.Qian is in part supported by the National Science Foundation Award CCF-1553281. We also thank Prof. Andrew Jiang and Junru Wu for helping experiments.

References