HASA: Hybrid Architecture Search with Aggregation Strategy for Echinococcosis Classification and Ovary Segmentation in Ultrasound Images

Traditional NAS methods mainly search over a discrete set of candidate architectures, which are very computationally intensive. The recent work (Liu et al., 2018b) proposed to solve the problem in a differentiable way. The computing architecture (i.e., cell) is represented as a directed acyclic graph. As shown in Fig. 6, an ordered sequence of N nodes ${x_{1},\dots,x_{N}}$ (latent representation, e.g. feature map) are connected by directed edges $(i,j)$ (i.e. operations $o(\cdot)\in O$). The NAS process boils down to find the optimal cell structure.

Formally, each internal node is computed as the sum of all its predecessors within a cell:

In each cell, there are candidate operations between two different nodes, including none, identity, max-pooling, separable convolution, and dilation convolution. To further compress the model and improve the performance, we introduce the MixConv layer and SE block (see Fig. 4) into the candidate operations.

For example, a MixConv35 layer combines the convolutional kernels of sizes $[3\times 3]$ and $[5\times 5]$ to fuse the features extracted from different receptive fields. SE block exploits adaptive channel-wise feature recalibration to boost feature learning and model generalization ability. After the searching process, the two candidate operations with the highest probability value are selected for each edge.

Fig. 5 presents the configuration of a classification cell and a segmentation cell obtained after the searching process. They both share the same candidate operations. For different tasks, the searched cell structures have significant divergence. It can be observed that the classification cells are more likely to use MixConv as well as larger kernels, since learning scale-invariant features is more important for the classification task. In contrast, the segmentation cell prefers the SE operation. The recent work (Rundo et al., 2019) indicates that the SE block’s adaptive feature recalibration provides excellent cross-dataset generalization, thus boosting the segmentation performance. Considering our dataset is collected from different equipment, such searched results are reasonable.

The parameters in our framework consist of the network parameters $\alpha$ and weight of architecture $\omega$. They are simultaneously learned through a bi-level optimization strategy:

where $L_{train}$ is the training loss and $L_{val}$ is the validation loss. We optimize the weight of the architecture $\omega$ in the inner level and the network parameters $\alpha$ in the outer level respectively. The final architecture is then derived by:

Where $o^{(i,j)}$ denotes the operation of transformation from $node^{i}$ to $node^{j}$. $O=\{o^{(i,j)}\}$ is the set of candidate operations. It can be seen from Eq. 2 that the optimization of the network parameters $\alpha$ and the weight of architecture $\omega$ is intertwined and requires a good balance.

As illustrated in Fig. 6, we propose a progressive growing strategy to address this problem. As the searching process going on, the number of the stacked cells increases hierarchically while the candidate operations are dropped one-by-one. In other words, there are a lot of operations in the early stage of the search, and we need to focus on the optimization of the architecture $\omega$. Then, we switch the computation cost from the architecture to the network parameters $\alpha$ by reducing the candidate operations and constructing a deeper network. Thus, the search process can focus more on the inner optimization of the network in the deep stages. Additionally, modern neural networks are usually very deep (e.g. ResNet152(He et al., 2016)). Our proposed searching strategy is able to progressively build a deep network architecture, which is crucial for extracting discriminative features.

Specifically, the cell architecture searching is divided into $K$ stages. In each stage, we stack a new cell behind an old one (see Fig. 6). The newly stacked cell has the same network parameters as the old one. The operation that obtains the lowest score (i.e. highest validation loss) will be dropped stage-by-stage. The progressive growing process can be formulated as follows:

Where $k=1,2,\cdots,K$ indicates the stage index, and $\|O\|$ denotes the cardinal number of candidate operations set.

We first tested our method for the classification of hepatic echinococcosis. Echinococcosis is a zoonotic parasitosis caused by a larval infestation that mainly occur in the liver (McManus et al., 2003). The two main types of the disease are cystic echinococcosis (CE) and alveolar echinococcosis (AE). According to disease progression, CE can be classified into CL, CE1, CE2, CE3, CE4, and CE5, and AE is classified into AE1, AE2, and AE3 (see examples for different classes in Fig. 1). AE is caused by infection with echinococcus multilocularis and starts with infiltration type (AE1) which is characterized by a tumor-like, infiltrative, and destructive growth. AE2 is characterized by the occurrence of calcium salt deposition and calcification. Finally, the lesion proliferates into a huge lesion, and the central part of the lesion becomes hollow due to ischemic necrosis and liquefactive necrosis, which forms the liquefaction necrosis type (AE3). Details about the classification of CE can be found here (Group et al., 2003).

US is the preferred imaging modality for echinococcosis diagnosis (McManus et al., 2003); however, the diagnosis is non-trivial due to the complex sonographic appearance and high intra- and low inter-class variety (Fig. 1). It is a tough task which requires the strong discriminative power from the network design. The classification dataset consists of 9566 US images from 5028 patients who were diagnosed with hepatic echinococcosis. B-mode and convex probes were used. The image numbers for each class are as follows: AE1 673, AE2 261, AE3 183, CL 2050, CE1 1595, CE2 399, CE3 1499, CE4 1974, and CE5 932. Some patients had more than one lesion in the liver, and for some lesions more than one image was saved by the doctors. Due to these reasons, the image number is not the same as the patient number. The ground-truth class label of each lesion was provided by abdominal sonographers with at least 5-year experience and further reviewed by a senior expert with 21-year experience. The lesions were manually annotated by sonographers with bounding boxes. All lesions were cropped according to the bounding boxes and resized to 224$\times$224. The dataset was then divided into the training and test sets according to a ratio of 7$:$3.

We further tested our method on an ovary dataset for ovary and follicle segmentation. An ovary image contains multiple follicles. Accurately segmenting the ovary and follicles is non-trivial (Li et al., 2019). The main challenges for automatic segmentation include very large variations in the size of ovary and follicles, poor contrast of ovary against the background, and touching boundary among follicles. The ovary ultrasound images were obtained by transvaginal ultrasonography. Ten sonographers with at least 5-year experience manually delineated all the ovaries and follicles in each image, and the results were reviewed by a senior expert with 20-year experience in pelvic ultrasound. The dataset consists of 3204 images, in which 2509 images from 153 patients were selected as the training set. The rest were used as the test set. All images were resized to 448$\times$448.

To avoid potential bias, when splitting data we make sure that the images from the same patient went into either the training or test set. In addition, data augmentation strategies including normalization, rotation ($\pm$20^∘), and random horizontal flip were applied to the training images.

The pre-trained DarkNet53 was used as the classification backbone. DarkNet53 is the backbone structure of the famous detection model YOLO(Redmon et al., 2016). Inheriting the strengths of ResNet, DarkNet53 has excellent representation learning ability and can avoid the vanishing gradient problem when training a very deep neural network. The progressive growing strategy was used for stacking the cell architectures. Specifically, there were five initial cells at the beginning. In each stage, we stacked two cells behind. K was set to three. As a result, there were a total of 11 cells involved in the NAS search. The architecture search model was trained for 25 epochs in each stage. Cross entropy loss was used for the classification task. We set the batch size to 36 and used Adam optimizer with a learning rate of 0.025 to update the network parameters. The search time for finding our final classification architecture is about 1.5 hours, and the time for training the whole classification network is about 2.5 hours.

We adopted the pre-trained MixNet as the segmentation backbone. MixNet naturally integrates the dimensions of multiple convolution cores into a single convolution, which reduces the amount of convolution computation while potentially improving the accuracy and efficiency of the model. Several cell structures were stacked behind as a block, including a cell and bilinear interpolation at the beginning. There were a total of two blocks involved in the progressive growing strategy. In each stage, we stacked one cell in each block. The K was set to two. As a result, there were a total of six cells involved in the NAS search. Finally, we used a 4$\times$ up-sampling to generate the prediction directly, with the same resolution as the input. It is worth noting that the candidate operations do not contain max-pooling and average-pooling. Dice similarity coefficient (DSC) loss was used for the segmentation task. We set the batch size to 8 and used Adam optimizer with a learning rate of 0.0001 to optimize the model. The search time for finding our final segmentation architecture is about 9.5 hours, and the time for training the whole segmentation network is about 17.5 hours.

To verify the effectiveness of the proposed method for echinococcosis classification, we compared it with several other DNNs, including:

DarkNet53: DarkNet53 (Redmon & Farhadi, 2018) serves as a baseline for comparison since our method replaces the 4^th and 5^th stages of DarkNet53 with the searched cells. DarkNet53 mainly consists of convolutional layers and residual structures, which have powerful feature learning capabilities.

SPOS: SPOS (Guo et al., 2019) is an evolution learning based search method. it constructs a simplified supernet, where all architectures are single paths so that weight co-adaption problem is alleviated. All architectures are trained equally with a uniform path sampling strategy.

ProxylessNAS: Unlike proxy NAS algorithms that search for building blocks on proxy tasks such as training for fewer epoch or stating with a smaller dataset, proxylessNAS (Cai et al., 2018) directly learns the architectures on the target task without any proxy and significantly reduces the computational cost of architecture search.

DARTS: Instead of searching for architectures over a discrete and non-differentiable search space, DARTS (Liu et al., 2018b) relaxes the search space to be continuous, thus allowing efficient search of the architecture using gradient descent.

P-DARTS: Due to the large gap between the depths in search and evaluation scenarios, previous differentiable search methods suffer from lower accuracy in evaluating the searched architecture or transferring it to another dataset. P-DARTS (Chen et al., 2019) uses search space approximation and regularization to allow the depth of searched architectures to grow gradually during training.

Pruning: To test whether pruning approaches can increase the performance of DarkNet53 with a smaller number of parameters, we adopted the L1-norm based channel pruning approach (Li et al., 2016). For a fair comparison, we balanced the pruning process to ensure the pruned model had the same-level parameters and size as ours.

For a fair comparison, all NAS-based experiments used a similar training scheme. Specifically, the training data is further separated into a 30% subset for architectures search and the remaining 70% for updating the network parameters. Table 1 presents the accuracy, precision, recall, and F1-score of different models. Compared with the DarkNet53, the proposed HASA_Dc achieved better performance with less resource consumption. Specifically, our HASA_Dc reduced about 40$\%$ parameters compared with DarkNet53 while achieving $2\%$ higher F1-score. The pruning approach significantly reduces the parameters but slightly degrades the F1-score by 1.2 percent compared with the pre-trained DarkNet53 without pruning. The HASA_Dc also got much better performance compared with the models searching from scratch (i.e., SPOS, proxylessNAS, DARTS, and P-DARTS). These results proves that the proposed HAS framework is superior to the vanilla NAS models in the classification task. We conjecture that this is due to the better feature extraction ability of our hybrid model.

By default our HAS replaces the 4^th and 5^th stages with stacked cells. We also implemented several variants of our method to compare different architectures of the hybrid model:

HAS-stage5: We only replaced the 5^th stage with stacked cells. The channels of the cells were set to $512,1024$ for being consistent with DarkNet53.

HAS-w/o-MS: To demonstrate the advantage of introducing the MixConv and SE block into the candidate operations, we experimented with the HAS framework without the MixConv and SE block operations.

HASA_Dc: We added the Dense-cell to our HAS framework to show the effectiveness of this component

In Table 1, it is interesting to note that HAS performed better than HAS-stage5. This indicates that only replacing the layers close to the final classification layer is not enough. Meanwhile, the HAS had higher accuracy than HAS-w/o-MS, which is attributed to the introduction of the MixConv and SE block operations. After introducing the Dense-cell, the resulting HASA_Dc further improved the classification performance on all the metrics, which shows the efficacy of our re-aggregation strategy.

We then performed paired-sample t-test to see whether the difference of performance is significant between our HASA_Dc and other methods. Specifically, the probabilities of belonging to the ground-truth class predicted by our method and the other method were used as the scores for each sample. P-value less than 0.05 is considered statistically significant. As shown in Table 2, our proposed HASA_Dc performs statistically better than all other methods.

Moreover, we evaluated the effect of the number of searched cells used in the evaluation stage on the classification accuracy of HASA_Dc. Note that the number of cells used in the evaluation stage does not have to be the same as that in the search stage. As in previous studies on NAS (Chen et al., 2019; Liu et al., 2018b), the purpose of the search stage is to discover the best cell architecture. It can be observed in Table 3 that as the number of cells increased from 3 to 7, the accuracy of the model improved by 2.33%. However, further increasing the cell number did not improve the performance but required higher computational costs. This is consistent with the general knowledge that to some extent the performance of the model improves as the number of parameters increases. Therefore, we need to balance the size and performance of the model.

Fig. 8 shows the confusion matrices of the baseline DarkNet53 and the proposed HASA_Dc. It can be seen that the classification performance of HASA_Dc was improved compared with that of DarkNet53. Specifically, most off-diagonal values of HASA_Dc were smaller than those of DarkNet53 (examples are shown in the red, yellow, and green boxes), which indicates that fewer samples were misclassified by HASA_Dc.

To visualize the learned features by the two models, we used t-SNE to display the 2D embedding. Fig. 9 (left) is the embedding learned by the vanilla DarkNet53 network. It corresponds well with the clinical experience of doctors, where CL and CE1, CE4 and AE1, CE5 and AE2 are difficult to be differentiated as they are similar to each other. Fig. 9 (right) shows the embedding learned by the HASA_Dc. Note that the samples of the same class became closer and those challenging classes became more separable. This proves that the proposed HASA_Dc is more capable of disentangling the inter-class similarities, which is vital to the fine-grained classification.

We also show some example images along with the predicted labels by the pre-trained DarkNet53 and our method HASA_Dc in Fig. 10. The images in the first two rows all belong to CE3. These images were misclassified into CE1, CE2, or CE4 by the DarkNet53, which is consistent with the statistics in the confusion matrix in Fig. 8. In contrast, our HASA_Dc successfully classified them. There are also cases that both methods failed. In the third row, both DarkNet53 and HASA_Dc misclassified AE3 into AE1 or CE3. The misclassification of AE3 into AE1 (Fig. 10h) is probably due to the poor contrast of the lesion. The last row shows three examples of CE1 that were misclassified into CE3 or CL by both methods. As shown in the confusion matrix in Fig. 8, CL and CE1 are the most difficult classes to be differentiated. They are both unilocular, cystic lesions with uniform anechoic content. Their difference is that for CE1 the cyst wall is visible while for CL it is not. However, the two methods may not well capture this very subtle difference.

We compared our method with the the following methods to verify its effectiveness in the segmentation task:

DeepLabv3: DeepLabv3 (Chen et al., 2017) addresses the problem of segmenting objects at multiple scales. It consists of modules that use atrous convolution in cascade or in parallel to capture multi-scale context by adopting multiple atrous rates. In the original DeepLabv3, Xception was used as the backbone. In our study, since we used MixNet as the backbone in our method, we also tested DeepLabv3 with the MixNet as the backbone and referred to this setting as the baseline.

U-Net: Based on fully convolutional network, U-Net (Ronneberger et al., 2015) modifies and extends this architecture to enable training with fewer images and yielding more precise segmentation. It consists of a contracting path to capture context and a symmetric expanding path to enable precise localization.

Gated-SCNN: Instead of processing color, shape, and texture information all together in a deep network, Gated-SCNN (Takikawa et al., 2019) is a two-stream architecture for semantic segmentation that explicitly wires shape information as a separate processing branch that processes information in parallel to the classical stream. Higher-level activations in the classical stream are used to gate the lower-level activations in the shape stream, which can remove noise and let the shape stream focus on processing the boundary-related information.

Pruning: To test whether the segmentation performance of the baseline (i.e., DeepLabv3 using MixNet as backbone) can be improved with a smaller number of parameters using the L1-norm based channel pruning approach (Li et al., 2016), we controlled the pruning process to ensure the pruned model had the same-level parameters and size as our method for a fair comparison.

SPOS, ProxylessNAS, and DARTS: Similar to the proposed hybrid framework, the searched architectures by the three NAS algorithms (Guo et al., 2019; Cai et al., 2018; Liu et al., 2018b) were respectively combined with the MixNet backbone.

HASA_Ac: HASA_Ac represents our HAS framework with ASPP-cell. This is to verify the effectiveness of the cell re-aggregation strategy.

The baseline is the DeepLabv3 using MixNet as the backbone. For consistency, all NAS-based experiments shared a similar training scheme. The training data is further separated into a 30% subset for architecture search, and the remaining 70% is used to update the network parameters. As shown in Table 4, compared with those state-of-the-art architectures, our proposed method can achieve not only higher performance but also better parameter efficiency. Specifically, for follicle and ovary, the mean DSC of the HASA_Ac is improved by 0.59% and 0.51% compared with the baseline, respectively. The improvement of the mean JC is 0.62% and 0.77%, respectively. This indicates that the segmentation results of the HASA_Ac are closer to the ground truth. For HD and ASD which are the smaller the better, the proposed HASA_Ac also gets lower errors than the baseline, reducing 1.41 $pixels$, 0.15 $pixels$, and 2.09 $pixels$, 0.27 $pixels$ for follicle and ovary segmentation respectively, which shows that our HASA_Ac can achieve better boundary estimation. When compared with the advanced NAS methods, including SPOS, ProxlessNAS, and DARTS, both HAS and HASA_Ac present better performance on all metrics with the least parameters.

Although applying the pruning approach to the baseline halved the number of the parameters, the performance became worse than that of the baseline. In addition, we compared our method with a gated shape CNN architecture (i.e., Gated-SCNN) (Takikawa et al., 2019), The results in Table 4 show that our proposed method outperforms the Gated-SCNN in all evaluation metrics for segmenting both ovary and follicles. It is also noted that our method has far fewer parameters than the heavy Gated-SCNN.

Fig. 11 qualitatively compares the results of the proposed HASA_Ac and the baseline (DeepLabv3 using MixNet as the backbone) on the ovary and follicle segmentation tasks. The images in the first three rows are good examples for both methods. It can be observed that the segmentation results of the HASA_Ac are closer to the ground truth than the baseline. Particularly, HASA_Ac shows obvious advantages over the baseline in recognizing the tiny follicles (see the second and third rows). The last three rows show three modest segmentation examples. The misrecognition is mainly attributed to the low contrast between the region-of-interest structures and the surrounding tissue. For the image in the last row, both methods tended to misrecognize the hypoechoic areas outside the ovary as follicles due to the very weak signals of boundary.

It is worth noting that the proposed HASA costs fewer parameters while improving the segmentation performance. Equipped with the progressive growing search, our HAS can reduce about 52% parameters when compared with the baseline. This may indicate that our method would be more efficient and economic for the devices with limited computation resources, such as laptop US devices. By further introducing the cell re-aggregation strategy, i.e., ASPP-cell, our HASA can obtain further improvement on all metrics for both the follicle and ovary without adding any extra computation burden.