A Global Self-Attention Memristive Neural Network for Image Restoration

Abstract

Recently, using the idea of non-local operations, various non-local networks and the Vision Transformer have been proposed to model the long-range pixel dependencies, addressing the limitation of Convolutional neural networks(CNNs). However, most of these models cannot adaptively process images with different resolutions, and their large number of parameters and computational complexity make them unfavorable for edge devices. In this paper, we propose an efficient Global Self-Attention Memristive Neural Network (GSA-MNN) for image restoration and present a memristive circuits implementation scheme for GSA-MNN. GSA-MNN can both extract global and local information from images, which can be flexibly applied to different resolution images. Specifically, the Global Spatial Attention Module (GSAM) and the Global Channel Attention Module (GCAM) are designed to complete the modeling and inference of global relations. The GSAM is used to model global spatial relations between the pixels of the feature maps, while the GCAM explores global relations across the channels. Moreover, a multi-scale local information extraction module is proposed to deal with image regions with complex textures. Furthermore, we provide a modular designed circuit implementation scheme for these three modules and the entire GSA-MNN. Benefiting from the programmability of the memristor crossbars, three kinds of image restoration tasks: image deraining, low-light image enhancement, and image dehazing are realized on the same circuit framework by adjusting the configuration parameters. Experimental comparisons with over 20 state-of-the-art methods on 10 public datasets show that our proposed GSA-MNN has superiority.