PaperlayerView exact revision on arXiv

Source-linked AI summary

C2FNAS: Coarse-to-Fine Neural Architecture Search for 3D Medical Image Segmentation

Qihang Yu, Dong Yang, Holger Roth, Yutong Bai, Yixiao Zhang, Alan L. Yuille, Daguang Xu

arXiv:1912.09628v2cs.CV

TL;DR

Designing effective 3D segmentation networks across task contexts is difficult, and NAS can suffer from search-to-deployment inconsistency under memory constraints. C2FNAS searches topology first and cell operations second, then applies the pancreas-searched network across ten MSD tasks, reporting state-of-the-art performance and transferability.

  • Problem

    Choosing or designing 3D segmentation networks for different tasks is difficult, while NAS can have network- or input-size inconsistencies between search and deployment under 3D memory constraints.

    Source: §1 · p. 1

  • Method

    C2FNAS searches macro-level topology in a coarse stage and then selects micro-level cell operations in a fine stage to build a 3D segmentation network from scratch.

    Source: §3.2 · pp. 3–4

  • Results

    Across 10 MSD tasks, the consistent C2FNAS model outperforms nnU-Net and modified or ensembled 2D and 3D U-Nets, with state-of-the-art performance reported using one pancreas-searched network.

    Source: §1 · p. 2, §6 · p. 8

  • Takeaways & Limitations

    The authors report that the pancreas-searched network and shared training/testing hyperparameters demonstrate robustness and transferability across the other MSD tasks.

    Source: §1 · p. 2

    The search is computationally expensive: coarse search takes five days on 64 V100 GPUs, and fine-stage training and search take additional GPU time.

    Source: §4.1 · p. 7

Abstract

from arXiv · show

3D convolution neural networks (CNN) have been proved very successful in parsing organs or tumours in 3D medical images, but it remains sophisticated and time-consuming to choose or design proper 3D networks given different task contexts. Recently, Neural Architecture Search (NAS) is proposed to solve this problem by searching for the best network architecture automatically. However, the inconsistency between search stage and deployment stage often exists in NAS algorithms due to memory constraints and large search space, which could become more serious when applying NAS to some memory and time consuming tasks, such as 3D medical image segmentation. In this paper, we propose coarse-to-fine neural architecture search (C2FNAS) to automatically search a 3D segmentation network from scratch without inconsistency on network size or input size. Specifically, we divide the search procedure into two stages: 1) the coarse stage, where we search the macro-level topology of the network, i.e. how each convolution module is connected to other modules; 2) the fine stage, where we search at micro-level for operations in each cell based on previous searched macro-level topology. The coarse-to-fine manner divides the search procedure into two consecutive stages and meanwhile resolves the inconsistency. We evaluate our method on 10 public datasets from Medical Segmentation Decalthon (MSD) challenge, and achieve state-of-the-art performance with the network searched using one dataset, which demonstrates the effectiveness and generalization of our searched models.

1. Introduction

C2FNAS addresses the difficulty of designing 3D segmentation networks and the search-to-deployment inconsistencies that constrain existing NAS approaches. It searches network topology and cell operations in two stages, then evaluates a pancreas-searched model across ten MSD tasks.

1. Introduction

  • 3D segmentation is difficult because target structures can be tiny or vary in location, shape, and appearance, while abnormalities and anisotropic voxel spacing add challenges.The introduction illustrates these challenges with examples from MSD tasks.Source: §1 · p. 1, Figure 1 · p. 1
  • The authors motivate NAS because manually designing high-performance 3D networks requires expertise, and individually designed networks may work well only for particular tasks.They note that most researchers build on existing 3D U-Net or V-Net architectures with moderate modifications.Source: §1 · p. 1
  • Existing NAS can mismatch search and deployment network sizes, while topology search for 3D tasks faces memory constraints that cannot be eased by further shrinking already-small batches or crops.Reducing batch or crop size further can worsen performance or prevent convergence.Source: §1 · p. 2
  • C2FNAS searches macro-level network topology in a coarse stage, then selects micro-level cell operations in a fine stage.The coarse search is constrained by U-shape and encoder–decoder skip-connection priors; the fine stage considers 2D, 3D, and P3D operations.Source: §1 · p. 2
  • On 10 MSD tasks, C2FNAS achieves state-of-the-art performance, surpassing prior state of the art by 1% on pancreas and 2% on pancreas tumours.The network was searched using the pancreas dataset; the same model and training/testing hyperparameters were applied across the other tasks.Source: §1 · p. 2
  • The paper presents C2FNAS as a from-scratch search that avoids sacrificing network or input size and reports robustness and transferability across MSD datasets.Its stated contributions include a medical-segmentation-informed search space and method for each stage.Source: §1 · p. 2

2. Related Work

Related work reviews medical segmentation architectures and NAS methods, tracing both established U-Net variants and efforts to automate architecture or training choices. It highlights the limited use of NAS in medical image segmentation and the range of approaches explored.

2.1. Medical Image Segmentation

  • U-Net-based architectures remain prominent in medical segmentation, with 3D, residual, attention, slice-based, hybrid, and fused 2D/3D variants among prior approaches.The related work identifies nnU-Net as the MSD challenge winner, using task-adapted combinations of 2D U-Net, 3D U-Net, and cascaded 3D U-Net.Source: §2.1 · pp. 2–3

2.2. Neural Architecture Search

  • NAS research includes reinforcement-learning, evolutionary, differentiable, and one-shot methods, with parameter sharing reducing search cost to several GPU days or hours.In evolutionary methods, architectures are mutated and evaluated by validation accuracy.Source: §2.2 · p. 3
  • Medical-segmentation NAS has explored 2D/3D/P3D operations, automated hyperparameter or augmentation tuning, and cell searches within predefined 3D U-Net topologies.The cited medical-segmentation NAS works represent several distinct search targets and strategies.Source: §2.2 · p. 3

3. Method

C2FNAS separates topology search from cell-operation search to avoid mismatches in network or input size between search and deployment. It searches macro-level connectivity first, then assigns operations on the fixed topology.

3.1. Inconsistency Problem

  • 3D segmentation makes reducing already-small crops infeasible, while NAS memory constraints can otherwise force search-stage network or input sizes to differ from deployment.Auto-DeepLab reduces batch and crop sizes to fit memory, but the passage says this is infeasible for these tasks and can destabilize training.Source: §3.1 · p. 3

3.2. Coarse-to-fine Neural Architecture Search

  • A = S × C splits architecture into topology DAG S and cell-operation choices C; C2FNAS searches topology first, then assigns operations to nodes.This coarse-to-fine separation is intended to resolve search/deployment inconsistency in network and input size.Source: §3.2 · pp. 3–4

3.3. Coarse Stage: Macro-level Search

  • Each coarse topology uses standard 3D convolutions and is trained from scratch; U-shaped and skip-connection priors prune the space, while code-distance clusters guide evolutionary sampling.Topology similarity is measured by Euclidean distance between topology codes, and clusters are ranked using sampled networks’ performance.Source: §3.3 · pp. 4–5

3.4. Fine Stage: Micro-level Search

  • With topology fixed, fine search replaces cell operations with 3D, P3D, or 2D choices, using uniformly sampled single-path one-shot training and validation-based candidate selection.Cells are kept simpler in exchange for more cells under 3D memory constraints; after super-network training, K randomly sampled candidates are ranked by validation performance.Source: §3.4 · p. 5

4. Experiments

C2FNAS was evaluated on the ten-task MSD benchmark using an architecture searched on the pancreas dataset. The reported model outperformed state-of-the-art methods on most tasks, especially challenging ones, while remaining lighter than most popular 3D models.

4. Experiments

  • The evaluation covers all 10 MSD segmentation tasks, with test-set comparisons measured by Dice-Sørensen coefficient and averages reported across tasks and targets.Table 1 also includes single-model and five-fold ensemble results; the ensemble fuses predictions by majority voting.Source: §4 · p. 5, Table 1 · p. 7, §4.2 · p. 7

4.1. Implementation Details

  • Coarse search used a 12-cell network and priors that reduced the architecture space from 2.9 × 10^4 to 9.24 × 10^2.Cells receive the previous cell's output and, in specified cases, a second input; an evolutionary procedure clusters topology candidates and returns the best of 50 evaluated networks.Source: §4.1 · pp. 5–6
  • Search used the MSD Pancreas Tumours training set, with isotropic 1.0 mm resampling and training patches of 96 × 96 × 96.Training used random rotations and flips, SGD, and a summed Dice and cross-entropy loss; DSC was evaluated by sliding-window inference.Source: §4.1 · p. 6
  • Fine search selected 2D, 3D, or P3D operations for each cell using a single-path one-shot super-network.After super-network training, 2,000 candidates were sampled and initialized from its weights; validation used stride 48 to speed searching.Source: §4.1 · pp. 6–7

4.2. Segmentation Results

  • C2FNAS outperformed state-of-the-art methods on most MSD tasks, especially challenging ones, and was lighter than most popular 3D models.The authors also report a higher average across tasks and targets; unlike nnU-Net, C2FNAS used one model and hyperparameter setting across tasks.Source: Table 1 · p. 7, §4.2 · p. 7, Table 2 · p. 7

4.1. Implementation Details

  • The coarse search took 5 days on 64 V100 GPUs, while fine-stage super-network training and candidate search took 10 hours and 1 day, respectively, on 8 GPUs.The authors attribute the substantial search cost mainly to the time required to train and evaluate 3D models.Source: §4.1 · p. 7
  • Deployment retained the searched topology and operations under the same training setting, with patch-size adjustments for the small axial extent of Prostate and the small Hippocampus volumes.Prostate used 128 × 128 × 32 patches; Hippocampus data were up-sampled to 96 × 96 × 96.Source: §4.1 · p. 7

5. Ablation Study

The ablations examine the effects of fine-stage operation search, proxy-dataset choice, and adding model scaling. Fine search improved results on pancreas and lung, while scaling offered a performance-capacity trade-off that the authors excluded from the main experiments.

5.1. Coarse Stage versus Fine Stage

  • Fine-stage search improved performance on the pancreas target and also increased generality, yielding better performance on lung.The comparison used five-fold cross-validation for networks from coarse and fine stages.Source: §5.1 · p. 8

5.2. Search on Different Datasets

  • Pancreas-searched and lung-searched topologies each performed better on their own task, yet both performed well on the other dataset.The lung proxy dataset contained 64 cases, compared with 282 for pancreas.Source: §5.2 · p. 8

5.3. Incorporate Model Scaling as Third Stage

  • Model scaling increased capacity and performance but also raised parameter counts and FLOPs substantially, so the authors omitted it from the main experiment to keep the model moderate in size.The study scaled filter numbers using multipliers from 0.25 to 2.0 and reported single-fold validation results for pancreas and lung.Source: §5.3 · p. 8, Table 4 · p. 8

6. Conclusions

The authors propose coarse-to-fine NAS to design a transferable 3D medical segmentation network. They report that it outperforms nnU-Net and modified or ensembled 2D and 3D U-Nets, without attention or pyramid modules.

6. Conclusions

  • The searched network outperforms MSD champion nnU-Net and a series of well-modified and/or ensembled 2D and 3D U-Nets.The authors describe it as using consistent model and hyper-parameters for all tasks.Source: §6 · p. 8
  • Coarse-to-fine NAS automatically designs a transferable 3D segmentation network for medical image segmentation.The authors position it for a setting where existing NAS methods cannot work well because 3D segmentation is memory-consuming.Source: §6 · p. 8
  • The network uses no attention or pyramid module, which the authors say makes it a more powerful 3D backbone than current popular architectures.Source: §6 · p. 8

Related papers

Loading 1912.09628v2…