--- a +++ b/EEGLearn/model.py @@ -0,0 +1,224 @@ +##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ +## Created by: Yang Wang +## School of Automation, Huazhong University of Science & Technology (HUST) +## wangyang_sky@hust.edu.cn +## Copyright (c) 2018 +## +## This source code is licensed under the MIT-style license found in the +## LICENSE file in the root directory of this source tree +##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ + +#coding:utf-8 + +import tensorflow as tf + +def my_conv2d(inputs, filters, kernel_size, strides=(1, 1), padding='same', activation=None, name=None, reuse=None): + return tf.layers.conv2d(inputs=inputs, filters=filters, kernel_size=kernel_size, strides=strides, padding=padding, activation=activation, + kernel_initializer=tf.truncated_normal_initializer(stddev=0.1), bias_initializer=tf.constant_initializer(0.1), name=name, reuse=reuse) + +def build_cnn(input_image=None, image_size=32, n_colors=3, activation_function=tf.nn.relu, reuse=None, name='VGG_NET_CNN'): + # VGG_NET 32 # [samples, W, H, colors] + with tf.variable_scope(name, reuse=reuse): + input_image = tf.reshape(input_image, shape=[-1, image_size, image_size, n_colors], name='Reshape_inputs') + # layer_1 # 4个3*3*32 + + h_conv1_1 = my_conv2d(input_image, filters=32, kernel_size=(3,3), activation=activation_function, name='conv1_1') + h_conv1_2 = my_conv2d(h_conv1_1, filters=32, kernel_size=(3,3), activation=activation_function, name='conv1_2') + h_conv1_3 = my_conv2d(h_conv1_2, filters=32, kernel_size=(3,3), activation=activation_function, name='conv1_3') + h_conv1_4 = my_conv2d(h_conv1_3, filters=32, kernel_size=(3,3), activation=activation_function, name='conv1_4') + h_pool1 = tf.layers.max_pooling2d(h_conv1_4, pool_size=(2,2), strides=(2,2), padding='same', name='max_pooling_1') # shape is (None, 16, 16, 32) + + # layer_2 + h_conv2_1 = my_conv2d(h_pool1, filters=64, kernel_size=(3,3), activation=activation_function, name='conv2_1') + h_conv2_2 = my_conv2d(h_conv2_1, filters=64, kernel_size=(3,3), activation=activation_function, name='conv2_2') + h_pool2 = tf.layers.max_pooling2d(h_conv2_2, pool_size=(2,2), strides=(2,2), padding='same', name='max_pooling_2') # shape is (None, 8, 8, 64) + + # layer_3 + h_conv3_1 = my_conv2d(h_pool2, filters=128, kernel_size=(3,3), activation=activation_function, name='conv3_1') + h_pool3 = tf.layers.max_pooling2d(h_conv3_1, pool_size=(2,2), strides=(2,2), padding='same', name='max_pooling_3') # shape is (None, 4, 4, 128) + + return h_pool3 + + +def build_convpool_max(input_image, nb_classes, image_size=32, n_colors=3, + n_timewin=7, dropout_rate=0.5, name='CNN_Max', train=True, reuse=False): + """ + Builds the complete network with maxpooling layer in time. + + :param input_image: list of EEG images (one image per time window) + :param nb_classes: number of classes + :param image_size: size of the input image (assumes a square input) + :param n_colors: number of color channels in the image + :param n_timewin: number of time windows in the snippet + :return: a pointer to the output of last layer + """ + with tf.name_scope(name): + with tf.name_scope('Parallel_CNNs'): + convnets = [] + # Build 7 parallel CNNs with shared weights + for i in range(n_timewin): + if i==0: + convnet = build_cnn(input_image[i],image_size=image_size,n_colors=n_colors, reuse=reuse) + else: + convnet = build_cnn(input_image[i],image_size=image_size,n_colors=n_colors, reuse=True) + convnets.append(convnet) # list contains [None, 4, 4, 128] + convnets = tf.stack(convnets) # [n_timewin, nSamples, 4, 4, 128] + convnets = tf.transpose(convnets, [1,0,2,3,4]) # [nSamples, n_timewin, 4, 4, 128] + + with tf.variable_scope('Max_pooling_over_flames'): + # convpooling using Max pooling over frames + convnets = tf.reshape(convnets, shape=[ -1, n_timewin, 4*4*128, 1]) + convpool = tf.nn.max_pool(convnets, # [nSamples, 1,4*4*128, 1] + ksize=[1, n_timewin, 1, 1], strides=[1, 1, 1, 1], padding='VALID', name='convpool_max') + + + convpool_flat = tf.reshape(convpool, [-1, 4*4*128]) + h_fc1_drop1 = tf.layers.dropout(convpool_flat, rate=dropout_rate, training=train, name='dropout_1') + # input shape [batch, 4*4*128] output shape [batch, 512] + h_fc1 = tf.layers.dense(h_fc1_drop1, 512, activation=tf.nn.relu, name='fc_relu_512') + # dropout + h_fc1_drop2 = tf.layers.dropout(h_fc1, rate=dropout_rate, training=train, name='dropout_2') + # inputshape [batch, 512] output shape [batch, nb_classes] # the loss function contains the softmax activation + prediction = tf.layers.dense(h_fc1_drop2, nb_classes, name='fc_softmax') + + return prediction + +def build_convpool_conv1d(input_image, nb_classes, image_size=32, n_colors=3, + n_timewin=7, dropout_rate=0.5, name='CNN_Conv1d', train=True, reuse=False): + """ + Builds the complete network with 1D-conv layer to integrate time from sequences of EEG images. + + :param input_image: list of EEG images (one image per time window) + :param nb_classes: number of classes + :param image_size: size of the input image (assumes a square input)S + :param n_colors: number of color channels in the image + :param n_timewin: number of time windows in the snippet + :return: a pointer to the output of last layer + """ + with tf.name_scope(name): + with tf.name_scope('Parallel_CNNs'): + convnets = [] + # Build 7 parallel CNNs with shared weights + for i in range(n_timewin): + if i==0: + convnet = build_cnn(input_image[i],image_size=image_size,n_colors=n_colors, reuse=reuse) + else: + convnet = build_cnn(input_image[i],image_size=image_size,n_colors=n_colors, reuse=True) + convnets.append(convnet) + convnets = tf.stack(convnets) + convnets = tf.transpose(convnets, [1,0,2,3,4]) + + with tf.variable_scope('Conv1d_over_flames'): + convnets = tf.reshape(convnets, shape=[ -1, n_timewin, 4*4*128, 1]) + convpool = my_conv2d(convnets, filters=64, kernel_size=(3, 4*4*128), strides=(1, 1), padding='valid', activation=tf.nn.relu, name='convpool_conv1d') + + + with tf.variable_scope('Output_layers'): + convpool_flat = tf.reshape(convpool, [-1, (n_timewin-2)*64]) + h_fc1_drop1 = tf.layers.dropout(convpool_flat, rate=dropout_rate, training=train, name='dropout_1') + h_fc1 = tf.layers.dense(h_fc1_drop1, 256, activation=tf.nn.relu, name='fc_relu_256') + h_fc1_drop2 = tf.layers.dropout(h_fc1, rate=dropout_rate, training=train, name='dropout_2') + prediction = tf.layers.dense(h_fc1_drop2, nb_classes, name='fc_softmax') + + return prediction + + +def build_convpool_lstm(input_image, nb_classes, grad_clip=110, image_size=32, n_colors=3, + n_timewin=7, dropout_rate=0.5, num_units=128, batch_size=32, name='CNN_LSTM', train=True, reuse=False): + """ + Builds the complete network with LSTM layer to integrate time from sequences of EEG images. + + :param input_image: list of EEG images (one image per time window) + :param nb_classes: number of classes + :param grad_clip: the gradient messages are clipped to the given value during + the backward pass. + :param image_size: size of the input image (assumes a square input) + :param n_colors: number of color channels in the image + :param n_timewin: number of time windows in the snippet + :param num_units: number of units in the LSTMCell + :return: a pointer to the output of last layer + """ + with tf.name_scope(name): + with tf.name_scope('Parallel_CNNs'): + convnets = [] + # Build 7 parallel CNNs with shared weights + for i in range(n_timewin): + if i==0: + convnet = build_cnn(input_image[i],image_size=image_size,n_colors=n_colors, reuse=reuse) + else: + convnet = build_cnn(input_image[i],image_size=image_size,n_colors=n_colors, reuse=True) + convnets.append(convnet) + convnets = tf.stack(convnets) + convnets = tf.transpose(convnets, [1,0,2,3,4]) # 调换轴 shape: (nSamples, n_timewin, 4, 4, 128) + + with tf.variable_scope('LSTM_layer'): + # (nSamples, n_timewin, 4, 4, 128) ==> (nSamples, n_timewin, 4*4*128) + convnets = tf.reshape(convnets, shape=[-1, n_timewin, 4*4*128], name='Reshape_for_lstm') + #lstm cell inputs:[batchs, time_steps, 4*4*128] + with tf.variable_scope('LSTM_Cell'): + lstm_cell = tf.contrib.rnn.BasicLSTMCell(num_units=num_units, forget_bias=1.0, state_is_tuple=True) + outputs, final_state = tf.nn.dynamic_rnn(lstm_cell, convnets, dtype=tf.float32, time_major=False) + # outputs.shape is (batch_size, time_steps, num_units) + outputs = tf.transpose(outputs, [1,0,2]) # (time_steps, batch_size, num_units) + outputs = outputs[-1] + + with tf.variable_scope('Output_layers'): + h_fc1_drop1 = tf.layers.dropout(outputs, rate=dropout_rate, training=train, name='dropout_1') + h_fc1 = tf.layers.dense(h_fc1_drop1, 256, activation=tf.nn.relu, name='fc_relu_256') + h_fc1_drop2 = tf.layers.dropout(h_fc1, rate=dropout_rate, training=train, name='dropout_2') + prediction = tf.layers.dense(h_fc1_drop2, nb_classes, name='fc_softmax') + + return prediction + + +def build_convpool_mix(input_image, nb_classes, grad_clip=110, image_size=32, n_colors=3, + n_timewin=7, dropout_rate=0.5, num_units=128, batch_size=32, name='CNN_Mix', train=True, reuse=False): + """ + Builds the complete network with LSTM and 1D-conv layers combined + + :param input_image: list of EEG images (one image per time window) + :param nb_classes: number of classes + :param grad_clip: the gradient messages are clipped to the given value during + the backward pass. + :param imsize: size of the input image (assumes a square input) + :param n_colors: number of color channels in the image + :param n_timewin: number of time windows in the snippet + :return: a pointer to the output of last layer + """ + with tf.name_scope(name): + with tf.name_scope('Parallel_CNNs'): + convnets = [] + # Build 7 parallel CNNs with shared weights + for i in range(n_timewin): + if i==0: + convnet = build_cnn(input_image[i],image_size=image_size,n_colors=n_colors, reuse=reuse) + else: + convnet = build_cnn(input_image[i],image_size=image_size,n_colors=n_colors, reuse=True) + convnets.append(convnet) + convnets = tf.stack(convnets) + convnets = tf.transpose(convnets, [1,0,2,3,4]) + + with tf.variable_scope('Conv1d_over_flames'): + convpool = tf.reshape(convnets, shape=[ -1, n_timewin, 4*4*128, 1]) + convpool = my_conv2d(convpool, filters=64, kernel_size=(3, 4*4*128), strides=(1, 1), padding='valid', activation=tf.nn.relu, name='convpool_conv1d') + conv1d_out = tf.reshape(convpool, [-1, (n_timewin-2)*64]) + + with tf.variable_scope('LSTM_layer'): + # (nSamples, n_timewin, 4, 4, 128) ==> (nSamples, n_timewin, 4*4*128) + convnets = tf.reshape(convnets, shape=[-1, n_timewin, 4*4*128], name='Reshape_for_lstm') + #lstm cell inputs:[batchs, time_steps, 4*4*128] + with tf.variable_scope('LSTM_Cell'): + lstm_cell = tf.contrib.rnn.BasicLSTMCell(num_units=num_units, forget_bias=1.0, state_is_tuple=True) + outputs, final_state = tf.nn.dynamic_rnn(lstm_cell, convnets, dtype=tf.float32, time_major=False) + # outputs.shape is (batch_size, time_steps, num_units) + outputs = tf.transpose(outputs, [1,0,2]) + lstm_out = outputs[-1] + + with tf.variable_scope('Output_layers'): + dense_in = tf.concat((conv1d_out, lstm_out), axis=1, name='concat_conv1d_lstm') # shape [batch, (n_timewin-2)*64+num_units] + h_fc1_drop1 = tf.layers.dropout(dense_in, rate=dropout_rate, training=train, name='dropout_1') + h_fc1 = tf.layers.dense(h_fc1_drop1, 512, activation=tf.nn.relu, name='fc_relu_512') + h_fc1_drop2 = tf.layers.dropout(h_fc1, rate=dropout_rate, training=train, name='dropout_2') + prediction = tf.layers.dense(h_fc1_drop2, nb_classes, name='fc_softmax') + + return prediction