keras中添加自己的层需要用到keras.layers.core.Lambda这一操作,并且lambda 属于一元操作符,用于接收keras的上一层的tensor,当添加的层比较简单时可以直接将操作列于lambda 后,比如
x1 = keras.layers.core.Lambda(lambda x: K.batch_dot(x, K.transpose(x))(x1)
其中,通过以上操作可为model添加一个求x和自身的转置做矩阵乘法的层layer,x1为上一层输出的tensor。
当操作比较复杂时,可定义一个函数来实现添加layer,比如构建一个名为PAM的函数,把所有复杂操作放于PAM函数即可,需要强调的是PAM只能是一元函数,道理很简单,因为tensorflow的流中间可以分叉,但总输入和总输出的flow肯定只有一支。需要注意的是,当在类class中进行模型定义时,PAM不用定义为类中的方法,而是应该放到类外,否则会报错:
TypeError: ('Not JSON Serializable:', <X_module_FPM_3att.X_module object at 0x7fe651821198>)
其中,X_module_FPM_3att是py文件名,X_module是定义模型的类名。
因此,采用函数定义层时用如下方式:
x1 = Lambda(lambda x: PAM(x))(x1)
def PAM(x):
m_batchsize, height, width, C = K.int_shape(x)
x1 = Conv2D(C // 8, (1, 1), padding='same')(x)
x2 = Conv2D(C // 8, (1, 1), padding='same')(x)
x3 = Conv2D(C , (1, 1), padding='same')(x)
x2 = K.permute_dimensions(x2, [0, 3, 1, 2])
proj_query = K.reshape(x1, [-1, width*height, C//8])
proj_key = K.reshape(x2, [-1, C//8, width*height])
energy = K.batch_dot(proj_query, proj_key)
attention = K.softmax(energy)
proj_value = K.permute_dimensions(x3,[0, 3, 1, 2])
proj_value = K.reshape(proj_value, [-1, C, width*height])
out = K.batch_dot(proj_value, K.permute_dimensions(attention, [0, 2, 1]))
out = K.reshape(out, [-1, C, height, width])
out = K.permute_dimensions(out, [0, 2, 3, 1])
out = keras.layers.add([x,out])
return out