NAS-ASDet: An Adaptive Design Method for Surface Defect Detection Network using Neural Architecture Search

The differentiable search method converts a discrete search space into a continuous differentiable form such that gradient descent can be applied to the search process to exceed the black-box optimization efficiency, such as reinforcement learning and evolutionary algorithms. The earliest gradient-based idea was proposed in DARTS [26]. But there remain two problems [27] in gradient-based DARTS: (a) Full candidate operations participate in the whole search process, which leads to longer search times and heavier computational overhead; (b) The transfer of rough decoupling may leads to performance variation. Many studies have improved DARTS: DATA [28] developed the ensemble Gumbel softmax estimator, which realized migration between the search and validation stages. PC-DARTS [29] proposed channel sampling and edge normalization technologies to reduce GPU resource consumption. P-DARTS [30] considered performance collapse and adopted a progressive search strategy. Even so, stably and efficiently searching strategies remain a popular research topic.

Auto-DeepLab [31] is generally recognized as the pioneering work demonstrating NAS application in image segmentation. It also extended the gradient optimization-based strategy to the segmentation task for the first time. NAS-Unet [14] improved on the baseline U-net using NAS and showed higher performance. FasterSeg [32] used a multibranch architecture to overcome model breakdown. DCNAS [33] designed a more complex supernet than Auto-DeepLab, which added cross-layer connections to the search space. DNAS [34] designed a three-level decoupled search strategy to reduce the training difficulty. In addition, in order to eliminate the deployment pressure of large networks, some methods [35, 36, 37, 38] also take into account the combined effect of NAS and hardware awareness to search for lighter architectures. However, most of the current NAS methods are designed based on natural scenes.

Although there have been some recent attempts of NAS in industrial scenarios, most of them are limited to defect classification or focused on specific tasks. For example, [39] and [40] realized defect classification of steel cracks using NAS. [41] explored an automated defect detection method for industry wood veneer, and [42] developed a NAS-based detection approach for analyzing defects in photovoltaic cells in electroluminescence images. Such methods do not meet the requirements of robustly designing precise defect localization in diverse industrial scenarios.

Therefore, the design method of NAS pixel-level defect detection network with high performance in industrial scenarios remains largely unexplored.

Since a small search space size should be used to reduce search difficulty, noting that many handcrafted architectures are based on repeated fixed blocks and inspired by [13], we desgin the search space for the lightweight network based on two types of repeatedly stacked lightweight cells: normal cells and reduction cells. Both of these cells follow the same construction pattern, wherein the output feature map of normal cell contains the same dimension as the input and half the size of reduction cell.

Many manually designed detection networks enhance their ability to detect defects in complex, irregular, and diverse contexts by using stronger multiscale representation. Therefore, inspired by Res2Net [44] to expand the receptive field, we design the cell to generate multiple available receptive fields at a finer cell level granularity. The cell-level search space is shown in Figure. 2(a).

As the industrial scene is more seriously affected by complex environments and lighting changes, the manual detection network usually adds attention mechanism to enhance the ability to extract important feature information, but when and how to add attention operation still requires researchers to analyze specific problems. Therefore, we propose to use attention operations to enrich the candidate operations, including channel attention [45] and spatial attention [46]. So that cells can adaptively focus on key defects in complex environments.

Since industrial scenarios have lightweight and low-computation requirements, we first collect basic operations that are widely used in detection networks (such as convolution, pooling, etc.), and then use separable convolutions instead of ordinary convolutions to limit cell lightweighting. And the final full candidate operation set $\mathcal{O}$ is summarized in Table 1.

Each cell is viewed as a directed acyclic graph consisting of multiple edges and $N$ nodes. Two of nodes ($C_{k-1}$,$C_{k-2}$) are inputs, one ($C_{k}$) is output, and the remaining $N-3$ are intermediate nodes. The transformation from the $i$-node to the $j$-node is connected by a mixed edge, which is a mixture of all candidate operations denoted $O^{(i,j)}$. To enable gradient optimization, the selection of discrete operations is relaxed as a continuous optimization problem. Specifically, we assign each candidate operation in each edge with an architecture weight $\alpha_{o}(\alpha_{o}\in\boldsymbol{\alpha})$, and $softmax$ is applied to ensure that $\alpha_{o}\in[0,1]$, $\sum\alpha_{o}=1$. The information flow between the $i$-node and $j$-node can be calculated:

where $i<j$ represents that $i$ is the forward node of $j$, and the data flow needs to be transmitted from $i$ to $j$. $x_{i}$ represents the $i$-th node feature map in cells, and each candidate operation $o(\cdot)$ belongs to candidate operation set space $\mathcal{O}$ with a relative weight $\alpha_{o}^{(i,j)}$. $\overline{o}^{(i,j)}(x_{i})$ represents the feature map corresponding to node $j$ after data is transmitted from node $i$ to node $j$.

Each intermediate node accepts all forward intermediate node feature maps and the output of the previous two cells:

where $n^{(i)}$ represents the feature map of the $i$-th intermediate node, $C_{m}$ denotes the forward cell output, and $m=k-1,k-2$ indicate the previous and previous-previous cells, respectively.

Then, the cell concatenates all intermediate nodes and the residual feature map sum as its output:

where $C_{k}$ represents the output of the $k$-th cell.

Finally, the architecture $\mathcal{A}$ with candidate operation set $\mathcal{O}$ can be formed by stacking multiple cells:

where $C_{nor}$ and $C_{red}$ denote the normal cell and reduction cell. $\{\overrightarrow{\cdot}\}$ represents the direction of information transmission.

This lightweight cell level search space provides several guarantees for detection performance: 1) Instead of using fixed connection and operations, the cell can be automatically searched and adaptively select an appropriate connection mode, such that more appropriate fine-grained features can be obtained and multiple available receptive fields can be provided, which improves the ability to detect defects. 2) Due to the addition of forward node $C_{k-2}$, the network can choose to accept more forward information and expand the receptive field. 3) Attention operations enrich the candidate operation set, and it is automatically determined by search, which enhances the ability to automatically focus on key defects in complex environments.

The segmentation NAS pioneered by Auto-DeepLab [31] usually uses single-level features for target prediction. However, the defect scales are varied, and single-layer features offer limited ability to recognize different scale targets. Noting that the effectiveness of combining multiscale features for performance improvement has been demonstrated in handcrafted architectures, we are motivated to integrate these effective knowledge of multiscale feature fusion into the proposed NAS framework to enhance the ability of the network to deal with multi-scale challenges, as shown in Figure. 2(b).

It should be noted that although handcrafted network performance proves that multiscale feature combination is effective, it does not indicate that in FPN-based detection architectures, balanced multiscale feature fusion is best. Different surface defects present different shapes and scales, which signifies that different level features have different levels of importance for detection, e.g., shallow features are more important for small-scale detection, and vice versa [7]. Therefore, we assign different weights at different feature levels, which are learned during the training process, toward adaptively adjusting the importance distribution of features and focusing on more useful information. The implementation details are as follows.

First, for given defect images, multilevel features are extracted by repeatedly stacking cells. The last layer features of each level $\{Fea_{1},Fea_{2},\cdots,Fea_{5}\}$ are used for subsequent adaptive feature fusion. To integrate feature maps on different scales, these maps are resized to $Fea_{1}$:

where $upsample^{n}(\cdot)$ indicates the upsampling operation by enlarging features by $n$ times, and $n\in\{1,2,4,8,16\}$. We concatenate $\overline{Fea_{i}}$ in the channel dimension to obtain the original feature fusion map $\mathcal{F}$:

Next, learnable weights are added to $\mathcal{F}$, which are used to adjust the importance distribution of different features. Specifically, we use the average pooling operation to obtain the representation vectors $V$ of each layer in $\mathcal{F}$. The fully connected layer $FC$ and nonlinear activation function $ReLU$ are used to encode $V$ to obtain $V^{\prime}$:

where $w_{\mathcal{F}}$ represents the parameter of $\mathcal{F}$.

Then, $V^{\prime}$ is normalized to $[0,1]$ through the $sigmoid$ function and used to represent feature importance:

Therefore, network can adaptively adjust the feature importance by learning $V_{f}$. The enhanced features can be expressed:

Finally, we resize the feature map in $\mathcal{F}_{enhance}$ using an upsampling operation to complete end-to-end segmentation prediction:

Thus, the lightweight network-level search space design is completed.

The network level search space considered multiscale features inspired by handcrafted architectures, which increases the lightweight cell-based search space expression. The generated architecture can not only extract data-driven multiscale features but also adaptively adjusts the feature’s importance distribution to cope with scale challenges, which guarantees effective detection under broadly diverse and complex factors.

We implement our network on PyCharm with the toolbox PyTorch. In the experiments, the method is trained and tested on NVIDIA Tesla A100 (with 40-GB GPU memory), CUDA vision of 11.4, Pytorch vision of 1.9.0, torchvision vision of 0.10.0, and CentOS Linux 8.0.

To evaluate the performance of NAS-ASdet, we conduct experiments on four industrial datasets that encompassed various challenges, including lighting conditions, scale variations, limited sample sizes, etc.

• 1.

  MCSD-C dataset: This dataset comes from multiple batches of motor commutator cylinder surface defects. Figure. 4 shows representative surface defect samples that appear on the commutator cylinder surface. The background environment of MCSD-C is complex and dynamic due to changes in production batches and line parameters. These factors leads to diverse lighting conditions and low-contrast defects, combined with defects exhibiting multi-scale variations, increase the difficulty of defect detection. We selected 566 defect samples with 256$\times$256. We used 445 images for training and the remaining samples for testing. Figure. 5 (1)-(2) rows shows example defect images and corresponding ground truth in MCSD-C.

  

• 2.

  RSDDs dataset[48]: RSDDs is a rail tread defect dataset collected from general/heavy transportation tracks. The detection challenges include severe surrounding noise interference and limited defect samples with complex contour shapes. RSDDs contains 128 defect samples with 55$\times$1250. We augment the original images with a 55-pixel sliding window, and finally obtain 191 training images and 82 testing images. Images are resized to 64$\times$64 during training. Figure. 5 (3)-(4) rows shows example defect images and corresponding ground truth in RSDDs.

• 3.

  KSDD dataset[49]: It is captured in a controlled industrial environment with visible surface cracks and contains only 54 defect samples. We augment the defect samples with a 500-pixel sliding window, resulting in 191 training images and 82testing images. KSDD primarily evaluates the capability for low-contrast defects with a limited number of sample images. Images are resized to 512$\times$512 during training. Figure. 5 (5)-(6) rows shows example defect images and corresponding ground truth in KSDD.

• 4.

  DAGM dataset[50]: The artificially created DAGM dataset provides a faithful representation of defects against a textured background. Its main challenge lies in the need to detect multiple types of defects with blurred boundaries under complex backgrounds and textures. In the original DAGM dataset, the defect regions are blanketed roughly by ellipses. In our experiment, four types of defects are selected and redefine the label at the pixel level. Finally we obtain 523 training images and 255 test images, with 512 × 512 original resolution. Figure. 5 (7)-(8) rows shows example defect images and corresponding ground truth in DAGM.

We compare NAS-ASDet with state-of-the-art handcrafted networks and existing NAS methods.

• 1.

  Handcrafted network: First, we select five classical handcraft networks designed for natural scenes as benchmarks for detection performance: FCN [17], U-Net[18], PSPNet[19], DeepLabV3+ [20], CSNet [51] (salient object detection technique). Secondly, we compare NAS-ASDet with four state-of-the-art manual defect detection networks, including defect segmentation frameworks and detection method based on salient object detection technique: PGA-Net[9],CSEPNet [22], TSERNet [21], and LSA-Net [52].

• 2.

  NAS-based method: We compare the NAS-ASDet performance with recent NAS segmentation methods: (a) Auto-DeepLab [31]: a macrosearch method for image segmentation, which is a classical baseline for NAS segmentation; (b) NAS-Unet[14]: which searches for specific cells and achieves good segmentation performance based on cell-based search space; (c) iNAS[35]: A hardware-aware saliency detection architecture with high performance based on NAS. (d) DNAS[34]: which uses the same search space as Auto-DeepLab, with a decoupled search framework to mitigate combinatorial explosion.

The above selected comparison methods cover image segmentation techniques, saliency object detection techniques, and lightweight network design, which helps to comprehensively evaluate the performance of our methods.

We use the same common metrics as [7] to evaluate model performance, including intersection over union (IoU), F1-Measure (F1), and pixel accuracy (PA). In addition, model Parameters (Params) and floating-point operation per second (FLOPs) are added because industrial sences are more sensitive to resource consumption. We use the total search time (Search_time) to measure the time consumption in NAS methods.

Given a defect detection dataset, the training process is divided into search and retraining steps.

• 1.

  Search Stage: This stage aims to generate deterministic detection architectures from a cell-based search space. The training set is divided 6:4 into an architecture training set and a weight training set, respectively. We define that each cell contains 4 intermediate nodes and 2 input nodes, and the network consists of 4 normal cells and 4 reduction cells. According to the search strategy described in Section 3.3, we establish 4 search stages, and the original candidate operations (Table 1) are gradually pruned in each stage according to $[top_{1}=7,top_{2}=4,top_{3}=2,top_{4}=1]$. The maximum epoch and batch size of each stage are set to 70 and 4, respectively. To avoid poor search performance caused by instability at the beginning of a search stage, we only use the weight training set to update network weights $w$ during the first 20 epochs of each stage. In the remaining epochs, architecture weights $\boldsymbol{\alpha}$ and network weights $w$ are updated by alternately using the architecture training set and weight training set. The Adam optimizer with a learning rate of 0.002, weight decay of 0.001 and momentum of 0.9 is used to update $\boldsymbol{\alpha}$, and the SGD optimizer with a learning rate of 0.005, weight decay of 0.0001 and momentum of 0.9 is used to update $w$.

• 2.

  Retraining Stage: After the search stage, a deterministic architecture is selected to replace the original supernet. The network weights $w$ are retrained for 500 epochs to mitigate training bias and ensure full network convergence. Multiscale features are fully considered according to the importance distribution. The other hyperparameters are set as in the search stage.

First, we begin the analysis with the performance of the designed network.

• 1.

  Comparison with handcrafted networks: As shown in Table 2, no fixed handcrafted network consistently achieves the best performance across the four different datasets (the best performance is highlighted by underline, including CSNet, CSEPNet, TSERNet, and LSANet), which demonstrates the necessity of using NAS technology. In contrast, our proposed NAS-based adaptive design method (NAS-ASDet) comprehensively outperforms the existing handcrafted networks and achieves state-of-the-art performance (IoU), both when compared to natural images and defect detection methods. It should be noted that unlike these manually designed networks that require empirical design, our approach adopts an automatic search-based design methodology, which only requires a few hours to automatically obtain the detection network on different datasets, significantly improving network design efficiency.

• 1.

  Comparison with NAS-based design methods: Our proposed NAS-ASDet consistently outperforms existing NAS-based methods. Particularly, on those datasets with more scarce samples (RSDD, KSDD, MCSD-C), compared with those NAS frameworks based on macro search (e.g. Auto-Deeplab, DNAS), our method shows more stable performance. We can observe an average IoU improvement of over 4% compared to the existing best-performing NAS methods, which shows the effectiveness of our refined search space. Moreover, we achieve this competitive result in relatively shorter search times. Except for DNAS, which completed the search faster on the MCSD-C dataset, our proposed method designed networks with optimal detection results in the shortest time. Compared to the difference in search performance of DNAS on MCSD-C, our method only takes an additional 33 minutes but achieved 6.44% IoU improvement. This demonstrates that NAS-ASDet outperforms existing NAS methods in terms of performance.

Second, we compare the computational complexity because lightweighting is crucial for industrial deployment. Table 2 presents the indicators of computational complexity, including Params and FLOPs. Across the four datasets, the networks’ size designed by NAS-ASDet are only about 1M-2M parameters with lower FLOPs. Compared to manually designed networks specifically for defect detection (PGANet, CSEPNet, TSERNet, and LSA-Net), NAS-ASDet utilizes only approximately 10% of the parameters and 5% of the FLOPs but achieving state-of-the-art performance. This achievement highlights the advantage of NAS-ASDet in industrial scenarios. Although NAS-ASDet has a slightly larger size than lightweight networks like CSNet and iNAS, the complexity increase is minimal (a maximum of 2M additional parameters and 5G additional FLOPs). While this increase does not significantly impose significant computational burden in industrial applications, but achieving an average IoU improvement of over 5%.

Conclusively, taking into account comprehensive performance metrics and computational complexity, NAS-ASDet surpasses other competitive methods and achieves state-of-the-art results.

In order to visually compare the differences in defect detection performance among different methods, we present the visual results of our proposed method and 13 comparative methods in Figure. 5. When those networks designed for natural images (Figure. 5 (d)-(h) columns) are applied to defect detection, the defect area only be roughly delineated, resulting in blurred details and missing some areas. Although those method designed specifically for surface defect detection (Figure. 5 (i)-(l) columns) can delineate the defects contour more comprehensively, they often exhibit false positives or false negatives in regions with low contrast features, leading to incorrect defect judgments. When facing the defects of scarce samples and low contrast, the existing NAS methods (Figure. 5 (m)-(p) columns) seem cannot well overcome the interference such as lighting and noise, which cannot show strong adaptability. In comparison, NAS-ASDet has more sensitive adaptability. As shown in Figure. 5 (1)-(2) rows, NAS-ASDet can capture low contrast defect features of different scales under different lighting conditions, without missing any defects. Even with a limited number of training samples and in the presence of environmental disturbances, our method can accurately distinguish the defect contour and obtain prediction results closer to the ground truth shown in Figure. 5 (3)-(4) rows. Figure (5)-(6) rows demonstrate that NAS-ASDet can accurately detect low contrast defects, even subtle cracks, with limited training samples. And it also shows strong adaptability by sensitively recognizing various types of defects with blurry boundaries in the presence of complex background textures, as shown in Figure. 5 (7)-(8) rows.

Therefore, our method can flexibly adapt to the challenge of defect detection, locate defects more accurately, and obtain better visual effects.

To provide a more intuitive representation of the searched network architecture, Figure. 6 illustrates the searched architectures on different datasets. As expected, the optimal cells obtained vary among these datasets, showcasing the ability of NAS-ASDet to automatically construct cells with diverse architectures driven by data. These architectures are further emphasized by the application of spatial attention and channel attention in different forms to the searched cell, allowing NAS-ASDet to adaptively adjust the feature map’s focus and enhance the cell’s representational power.

Although NAS-ASDet outperforms other competitive methods, there are still failure cases in some challenging situations. In Figure. 7, (a) and (b) columns illustrate typical defect regions with large proportions. However, NAS-ASDet may not consider the entire defect, resulting in partial loss when the defect area changes. For small defects with low contrast at the edge, as shown in the (c) and (d) columns of Figure. 7, our method may miss detection. In the (e), (f) column of Figure. 7, our method may overly focus on the vicinity of the defect area, leading to false detection in nearby regions. The main reason for this issue is the dataset availability and diversity, and we will address these deficiencies in our future work.

In NAS-ASDet, feature extraction is performed by repeatedly searching cells. We compare this approach with widely used classical feature extraction methods (VGG16 [11], ResNet50 [12], GoogLeNet [53], MobileNetV3 [54] and DenseNet [55]). During the comparison, we separately replace the feature extraction stage, while the other components of NAS-ASDet remain intact. As shown in Table 3, with the fixed connection mode, the extracted features are fixed and the performance is limited. Our search method can thus automatically determine a more suitable connection, making the network performance is better. Moreover, the searched architecture is more lightweight compared to these classical fixed feature extraction methods.

To consider the impact of fusing multiscale features, we evaluate different levels of feature maps ($Fea_{1},Fea_{2},Fea_{3},Fea_{4},Fea_{5}$) extracted by automatically searched cells. We also compare the performance of balanced fusion and adaptive fusion. As shown in Figure. 8: (1) The detection ability of fusing multiscale features is better than that of single-level features (Figure. 8(b) VS Figure. 8(a)), which proves that multiscale feature fusion can improve the detection performance; (2) Different single-level features show different performances (Figure. 8(a)). For example, $Fea_{2}$ achieves the best performance on RSDDs, while $Fea_{3}$ performs best on MCSD-C. This proves that the contributions of each level of features are different for different tasks, and it is necessary to integrate different scale features according to their importance so that the total performance can be enhanced. (3) The experiments show that adaptive fusion not only outperforms single-level features but also outperforms balanced fusion (Figure. 8(b)), proving that the adaptive adjustment of feature importance distribution is effective.