Searching Architecture And Precision For U-Net Based Image Restoration Tasks

Manually architecting Deep Neural Networks (DNNs) has led to the success of Deep Learning in many domains. However, recent DNNs designed using Neural Architecture Search (NAS) have exceeded manually designed architectures and have significantly reduced the human effort to develop complex networks. Current works use NAS to identify a cell architecture constrained by a fixed order of operations that is then replicated throughout the network. The constraints potentially limit the effectiveness of NAS in converging on a more efficient DNN architecture. In the first part of our paper, we propose Operation Search, a search on an enlarged topological space for U-net and its variants that retain efficiency. The idea is to allow for custom cells (operations and their sequence) at various levels of the network to maximize image quality while being sensitive to computation cost. In the second part of our paper, we propose custom quantization at various levels resulting in a mixed-precision network. Additionally, we increase the search efficiency by constraining the search space to use the same precision for both weights and activations at any level. This does not result in computational inefficiency because it matches the operand precisions supported by Tensor Core enabled GPUs.