A CNN identified by reinforcement learning-based optimization framework for EEG-based state evaluation
Yuxuan Yang, Zhongke Gao, Yanli Li, He Wang
Abstract
Abstract Objective. Electroencephalogram (EEG) data, as a kind of complex time-series, is one of the most widely-used information measurements for evaluating human psychophysiological states. Recently, numerous works applied deep learning techniques, especially the convolutional neural network (CNN), into EEG-based research. The design of the hyper-parameters of the CNN model has a great influence on the performance of the model. Therefore, automatically designing these hyper-parameters can save the time and labor of experts. This leads to the appearance of the neural architecture search technique. In this paper, we propose a reinforcement learning (RL)-based step-by-step framework to efficiently search for CNN models. Approach. Specifically, the deep Q network in RL is first used to determine the depth of convolutional layers and the connection modes among layers. Then particle swarm optimization algorithm is used to fine-tune the number and size of convolution kernels. Through this step-by-step strategy, the search space can be narrowed in each step for saving the overall time cost. This framework is employed for both EEG-based sleep stage classification and driver drowsiness evaluation tasks. Main results. The results show that compared with state-of-the-art methods, the high-performance CNN models identified by the proposed optimization framework, can achieve high overall accuracy and better root mean squared error in the two tasks. Significance. Therefore, the proposed optimization framework has a great potential to provide high-performance results for other kinds of classification and prediction tasks. In this way, it can greatly save researchers’ time cost and promote broader applications of CNNs.