Towards Cardiac Intervention Assistance: Hardware-aware Neural Architecture Exploration for Real-Time 3D Cardiac Cine MRI Segmentation

In this section, We will first introduce the network structure and 3D multi-scale neural architecture search space. Then we will demonstrate how we incorporate the latency estimated as an additional term into the loss function and jointly search for network architecture with high accuracy and low latency. Finally, we will discuss how the network is optimized and how to decode the discrete architecture once the search finishes. Figure 1 shows an overview of the 3D multi-scale neural architecture search space, which can be represented as a directed acylic graph (DAG). Each vertex $v_{i}$ represents a feature map. Each edge $e_{ij}$ represents an operation between two vertices $v_{i}$ and $v_{j}$. Following the recent cell-based neural architecture search strategy (Liu et al., 2018; Xu et al., 2019b; Weng et al., 2019), the best cell architecture is searched during the neural architecture search phase and then shared across the entire network. Note that a cell is defined as a network structure associated with each edge in the DAG.

Like existing cell-based NAS works (Weng et al., 2019; Liu et al., 2018), the structure of the cell is shared by the entire network, different cell structures can be defined according to the functionality of the cell. In our work, we define three types of cell architectures called non-scaling cell, contracting cell and expanding cell which are associated with the non-scaling edge, the contracting edge, and the expanding edge in the DAG in Figure 1 respectively. The non-scaling cell keeps the scale of the feature map. The contracting cell and the expanding cell are used to down-sample and up-sample the scale of the feature map.

Figure 2 shows the structure of the defined cells. The structure is shared among three types of cells discussed above. Each types of cell has different preprocessing block, which contains a scaling operation (e.g., maxpooling or upsampling). The dashed edges/lines inside the cell represent the primitive operations that need to be searched. The primitive operations will be discussed in section 3.3. Suppose there are $m$ intermediate nodes inside a cell. Then the total number of edges that need to be searched is $m(m+1)/2$. In the end, the outputs of all intermediate nodes are concatenated into a single feature map as the output of the cell.

The search space of our network contains cell level and architecture level. At the cell level, we can search for the best primitive operations between intermediate nodes as well as the edge connections among them. As shown in Figure 3 (a), suppose there are $3$ intermediate nodes in a cell, then the search space in the cell can be represented as a DAG consisting $5$ vertices, with $v_{0}$ be the output feature map of the preprocessing block and $v_{4}$ be the output feature map of the cell. Each edge represents a mixture of primitive operations. In our work, the search space for the primitive operations is normal 3D convolution, dilation 3D convolution, depth-wise-separable 3D convolution, 3D maxpooling, identity (skip connection) and zero (none operation) as shown in Table 1. Recent works (Liu et al., 2018; Cai et al., 2018) showed the feasibility and effectiveness of using continuous relaxation for an over-parameterized network that includes all candidate paths. We use the same strategy for our neural architecture search. Specifically, it can be seen in Figure 3 (b) that each edge is a mixture of $N$ candidate primitive operations. Let $OP=OP_{k}$ be one of the primitive operations from the $N$ candidate sets and $X_{j}$ be the feature map of node $v_{j}$. Then the feature map of $v_{j}$ can be calculated as:

where $\alpha_{i,j,k}$ is the weight of the corresponding primitive operation $k$ on the edge from $v_{i}$ to $v_{j}$, $p_{\alpha_{i,j,k}}$ is the probability of choosing this operation. $\gamma_{i,j}$ is the weight parameter of edge from $v_{i}$ to $v_{j}$. $p_{\gamma_{i,j}}$ represents the probability of this edge. During search, $p_{\alpha_{i,j,k}}$ is computed by using softmax on the weights of all candidate operations $\alpha_{i,j}$. $\gamma_{i,j}$ is computed by applying softmax on weights of all edges to coming to node $v_{i}$. After neural architecture search, we can prune the redundant operations and paths by choosing $\alpha_{i,j,k}$ and $\gamma_{i,j}$ with the highest probability.

Inspired by (Xu et al., 2019b), we further apply a partial channel connection strategy in the search scheme to reduce the memory cost during a search. Specifically, we use a hyper-parameter $k$ to split the channels of the feature map into two parts, $1/k$ part and $1-1/k$ part. Let $1/k$ part do the mixed primitive operations and the remaining part stay unchanged. Finally, the results of these two parts are concatenated to form the resulting feature. Therefore, we can rewrite equation 1 to:

in which $k_{i,j}$ is a sampling mask that samples $1/k$ portion of channels from tensor $X_{i}$. In Figure 4, we use the calculation of feature map $X_{3}$ in a cell to demonstrate the mechanism of partial channel connection. The idea of partial channel connection is to let part of the channels go through the operations and the rest part stay constant, thus reducing the computation.

As for the network level search, unlike Darts (Liu et al., 2018) and Nas-unet (Weng et al., 2019) which fix the network backbone during the search, our framework has a multi-scale search space that can identify various backbones, such as UNet, ResUnet, FCN, etc. In particular, in the network level, we hope to search for the connection between vertices in Figure 1 (not vertices in the cell) to find the optimal path (or sub-network) within the entire network search space. For each path from the input to output in the DAG, there is only one edge going into a vertex and one edge coming out of it. Neural architecture search aims to identify such path that can produce the best accuracy and latency. Since we are using the continuous relaxation search strategy, each edge is assigned with a weight parameter $\beta$, indicating the sampling probability of this connection in the graph. Let $X_{s}^{l}$ be the feature map in layer $l$. Scaling level $s$, $s\in\{0,1,2,3\}$. $Ep(*)$, $Ns(*)$ and $Ct(*)$ denote the expanding, non-scaling and contracting operations respectively. $\beta_{i,j}^{l-1}$ denotes the weight parameter of edge from $X_{i}^{l-1}$ to $X_{j}^{l}$, $\beta_{i}^{l-1}$ is the weight parameter of the skip connection from $X_{i}^{l-1}$ to $X_{i}^{l}$. Then the feature map $X_{s}^{l}$ can be calculated as:

Similar to cell structure parameters $\alpha$ and $\gamma$, softmax operation is also used to normalize real parameter $\beta$ to get the corresponding probability $p_{\beta}$. Note that $\beta$ is normalized based on the weights of all the edges that come out of a vertex. Some vertices may not contain certain edges. For example, the vertices in the first layer do not have expanding edges that come out of them.

Most existing NAS methods focus on reducing FLOPs. However, FLOPs may not always reflect the actual latency performance of a network (Wu et al., 2019; Cai et al., 2018). In order to simultaneously optimize model accuracy and efficiency, we apply the following loss function during search

where $w$ is the network weight parameters. $\alpha$, $\beta$ and $\gamma$ are architecture parameters. The first term $CE(w,\alpha,\beta,\gamma)$ is the cross-entropy loss for accuracy optimization. The second term $LAT(\alpha,\beta,\gamma)$ is the hardware-aware latency term aiming to optimize the latency performance of the searched network on the target device (GPU in our case). $\lambda$ is a hyper-parameter that controls the magnitude of the latency term. $\lambda$ is set to 0.0001 in all our experiments. Similar to FBNet (Wu et al., 2019), we use a latency lookup table to estimate the overall latency of the network based on the runtime of each primitive operation.

Prior hardware-aware NAS methods (Wu et al., 2019; Cai et al., 2018; Li et al., 2020) usually use the fixed backbone during a search. Therefore, their latency can be estimated by taking the summation of all operators under the assumption that the runtime of each operator is independent of others. This is true if the network search space only contains a single backbone structure. However, in our multi-scale framework, the data dependency is more complicated. The information on one vertex can choose different paths to go and those operations on different paths may run simultaneously using GPU parallelism. Therefore, simply sum the latency of all operations in the network may not reflect the actual latency of the searched network. To address this issue, we introduce a path-wise latency estimation method, in which the expected latency of the $n$ longest paths are estimated and used as the expected latency of the searched network. Intuitively, only the $n$ longest paths are kept in the final searched network and our goal is to minimize the latency of these paths. Like (Yan et al., 2020), $n$ is a hyperparameter that determines the trade-off between model complexity and efficiency, which will be discussed in section 3.5.

Expected latency at the cell level: Similar to previous work of latency estimation, we use a differentiable paradigm to solve the problem. In particular, the expected latency inside a cell is calculated by summing up the latency of preprocessing operation and all mixed operations among intermediate nodes. The expected latency of the mixed operation between vertices in a cell is calculated as

in which $lat_{s,s+1}^{l}[{MixedOP_{i,j}}]$ represents the latency of the mixed operation from $v_{i}$ to $v_{j}$ in a cell from layer $l$ to $l+1$ and scaling level $s$ to $s+1$. $GS(\cdot)$ is the Gumbel-Softmax sampling rule to make the architecture parameters differentiable to the latency loss term. $\alpha_{i,j}$ are the weight parameters of the primitive operations from $v_{i}$ to $v_{j}$. $\alpha_{i,j,k}$ is the weight parameter of the $k\-$th parameter from the candidate primitive operations. The indices of vertices in a cell are illustrated in Figure 3. $N$ is the total number of candidate primitive operations. $F(\cdot)$ is the latency estimating function that returns the latency of some operation $OP_{k}$. With the latency of mixed operation defined, we can then compute the expected latency of the cell as

where $lat_{s,s+1}^{l}$ is the latency of a cell from layer $l$ to $l+1$ and scaling level $s$ to $s+1$. $\gamma_{j}$ are the weight parameters of edges coming into vertex $v_{j}$, $\gamma_{j,t}$ is the weight from one of these edges.

Expected latency at the architecture level: Based on the expected cell latency estimated above, we define the latency of a path to be the summation of the cell latency of each layer in the path. This is based on the observation that the input feature of layer $l$ is only dependent on the features of the previous layer $l-1$. Specifically,

$\mathop{{}\mathbb{E}}[lat[path_{n}]]$ is the expected latency of top $n$-th longest path in the searched network. We will discuss the definition of path length and how to find the top $n$ longest path in Section 3.5. $L$ is the number of layers. $lat_{s,t}^{l}$ is the latency of all possible cells from layer $l$ to $l+1$ and scale $s$ to $t$. $\beta_{s,t}^{l}$ is the corresponding weight parameter. Then the expected latency of the entire network is simply the summation of latency of all top $n$ longest paths. It can be seen from Equation 7 to Equation 9 that the expected latency term is differentiable respect to architecture parameter $\alpha$, $\beta$, and $\gamma$. Therefore, by minimizing the latency loss, we are searching for the optimal architecture parameters with the best latency performance.

Similar to (Liu et al., 2018; Liu et al., 2019a), we formulate the architecture search as a continuous optimization problem. Specifically, we adopt the first-order approximation in (Liu et al., 2018) for searching. The training set is split into two parts $train_{weight}$ and $train_{arch}$ for weight parameter optimization and architecture parameter optimization respectively. The training contains two phases. For the first few epochs, the architecture parameters are frozen and only network weight parameter $w$ is updated by $\bigtriangledown_{w}\mathcal{L}_{train_{weight}}(w,\alpha,\beta,\gamma)$. Then, for the rest training epochs, network weight parameter $w$ and architecture parameters $\alpha,\beta,\gamma$ are updated in turn by $\bigtriangledown_{w}\mathcal{L}_{train_{weight}}(w,\alpha,\beta,\gamma)$ and $\bigtriangledown_{\alpha,\beta,\gamma}\mathcal{L}_{train_{arch}}(w,\alpha,\beta,\gamma)$.

Once the search finishes, we still need to decode the architecture parameters to final discrete architecture. At the cell level, we take the argmax on $\alpha$ to choose the most likely primitive operation. We also take the argmax on $\gamma$ to find the edge connection between vertices in a cell. Note that each vertex in a cell represents a feature map, and a latter vertex may connect to all of its previous vertices. Examples can be seen in Figure 3. $\gamma$ is the architecture parameter that controls these connections. In the final cell architecture, one vertex can only connect with one of its previous vertices.

Unlike the cell level decoding, at the architecture level, we cannot simply apply the same strategy on $\beta$ to obtain the final discrete architecture path. This is because ideally, we can only have one edge in each layer in a single path, and argmaxing on edges coming out of different vertices in the same layer makes no sense as they are not related. Note that the sum of weights of edges coming out of a vertex is always 1 and the probability of choosing the current edge is conditioned on the probability of previous edges. Therefore, we let each network vertex sample independently. Then the length of a path can be described as the product of all the edge weights in the path

where $p_{\beta^{l}}^{(n)}$ is the softmax probability of the edge’s parameter $\beta$ in layer $l$ of path $n$. Note that for different path $n$, $p_{\beta^{l}}^{(n)}$ may be different. Dynamic programming is used for finding the $n$-th longest path in the DAG.

In order to further improve segmentation performance, we use the same strategy as in (Yan et al., 2020). Instead of using the longest path as the final searched architecture, we fuse the top-$n$ longest paths in the DAG. Intuitively, larger $n$ represents a bigger network model and more computation resource, but higher model complexity and usually better segmentation accuracy. Therefore, $n$ is a hyperparameter that can be used to explore tradeoff between accuracy and hardware efficiency. Experiment results will show that a small level of path fusion (n=2) can improve latency while maintaining competitive accuracy.

We evaluate our neural architecture search framework using the ACDC MICCAI 2017 challenge dataset (Bernard et al., 2018) with additional labeling done by experienced radiologists (Wang et al., 2019). The task is to segment right ventricle (RV), myocardium (MYO) and left ventricle (LV) from 3D cardiac MRI cine for both end-diastolic and end-systolic phases instances. The dataset contains 150 exams from different patients with 100 for training and 50 for testing. Dice score and 5-fold cross-validation are used for evaluating the segmentation accuracy. we evaluate the test data by submitting the segmentation results of ED and ES instants to ACDC online evaluation platform (acdc challenge, 2017).

During the searching phase, we split the training set of ACDC by 1:1 into two parts: $train_{weight}$ and $train_{arch}$, each with 50 cases. $train_{weight}$ is used for network weight training and $train_{arch}$ is for architecture parameter training. After a search finishes, the searched architecture is then trained from scratch on the entire training set following the 5 fold cross-validation. We compare our work with the state-of-the-art NAS framework MS-NAS (Yan et al., 2020) for medical image segmentation. We modify MS-NAS to a 3D version to handle 3D MRI images. We also implement traditional 3D U-Net with different configurations (D4, IF32) and (D3, IF16), 2D U-net with different configurations (D5, IF64) and (D3, IF16). For example, (D4, IF32) represents the depth of the network is 4, the initial filter number is 32. 2D U-net is based on the implementation of (Wang et al., 2019).

In our 3D multi-scale NAS supernet, the number of layers is set to 8. The number of intermediate nodes in a cell is 3. We set $k=4$ for partial channel connection. The 3D cine MRI images are resized and cropped to 20$\times$128$\times$112. 3D cine MRI images are preprocessed using the same strategy as (Isensee et al., 2018) with the same data augmentation for a fair comparison. We optimize the network weight in the first 30 epochs. For the rest 30 epochs, network weight and network architecture parameters are optimized alternately. All these methods are fully trained with the same hyperparameters.

The evaluated networks and NAS frameworks are implemented using Pytorch. All experiments run on a machine with 16 cores of Intel Xeon E5-2620 v4 CPU, 256G memory, and four NVIDIA GeForce GTX 1080 GPUs.

Table 2 shows the comparison among 2D U-net, 3D U-net, NAS-Unet, MS-NAS and our proposed NAS framework on the ACDC 3D cardiac cine MRI dataset. From the table, we can see that MS-NAS 3D can achieve the best Dice score on all of the three anatomy structures. However, it cannot satisfy the timing constraint. On the other hand, our proposed approach can guarantee the timing constraint to be met (i.e., 3.5 times improvement on latency), while achieving competitive accuracy against MS-NAS, with at most of 0.8% accuracy loss or even without accuracy loss. Note that $n$ represents the number of fusion paths in the searched network. Larger $n$ means larger network, usually with better accuracy.

We have several detailed observations from this table: (1) For all competitors, only U-net 2D with the depth of 3 and the initial filter channels of 16 (denoted as U-net 2D-3-16) can satisfy the throughput constraint of ($\geq$22FPS), while our method can guarantee that the throughput meets the constraint. This is because i) for 3D medical images, the network FLOPs are usually very large. We have to carefully design the network structure in order to fulfill the real-time constraint; ii) MS-NAS uses much larger operators in the cell search space (e.g., a modified 3$\times$3 depth-wise-separable convolution which involves 4 convolution layers); iii) MS-NAS does not consider latency while searching. (2) The Dice score of U-net 2D-3-16 for LV, RV, and MYO are merely 0.767, 0.564, and 0.738, while ours can improve it to 0.911, 0.865, and 0.850, respectively. (3) Even compared with the results of U-net 3D-5-64 and U-net 3D-4-32 that cannot satisfy the real-time constraint, our method can achieve better accuracy. (4) With $n=1$, we can achieve better latency at the cost of lower Dice score. Specifically, compared to $n=2$, $n=1$ can achieve 6ms latency improvement. However, the average Dice score drops by about 0.9$\%$. This demonstrates that the proposed method can provide freedom for designers to make a better tradeoff on accuracy and hardware efficiency.

All the above results verify the importance of conducting hardware-aware NAS for real-time 3D cardiac cine MRI segmentation.

Figure 6 presents the searched cell structure (top) and network structure (bottom). In the identified network structure, two paths with the longest length are fused to form the final architecture. The search is a trade-off between accuracy and latency. It can be seen in the non-scaling cell that the network prefers weight-free operations (e.g, identity and maxpooling). This is because the non-scaling cell is less important than other cells (e.g., contracting cell) in segmentation networks.

There is an interesting observation on the second-longest path. It is shallower than the first path and upsampling is not done until the last layer. One potential reason is that NAS is trying to guarantee the timing performance of the searched network. Such a path is hard to identify manually, which is why we need NAS to help explore network structures.

Finally, Figure 5 shows the visualization of some segmentation examples at different timestamps in a cine MRI. We can observe that for various positions (e.g., base and apex) and different timestamps, our method can always accurately segment the targets. It can be seen that the RV segmentation results are not as accurate as LV and MYO. This is because in some slices the RV is very small and the boundary is unclear, which makes it hard to segment.