Self-Attention GAN in Keras - tensorflow

Question

I'm currently considering to implement the Self-Attention GAN in keras.
The way I'm thinking to implement is as follows:
def Attention(X, channels):
def hw_flatten(x):
return np.reshape(x, (x.shape[0], -1, x.shape[-1]))
f = Conv2D(channels//8, kernel_size=1, strides=1, padding='same')(X) # [bs, h, w, c']
g = Conv2D(channels//8, kernel_size=1, strides=1, padding='same')(X) # [bs, h, w, c']
h = Conv2D(channels, kernel_size=1, strides=1, padding='same')(X) # [bs, h, w, c]
# N = h * w
flatten_g = hw_flatten(g)
flatten_f = hw_flatten(f)
s = np.matmul(flatten_g, flatten_f.reshape((flatten_f.shape[0], flatten_f.shape[-1], -1))) # [bs, N, N]
beta = softmax(s, axis=-1) # attention map
flatten_h = hw_flatten(h) # [bs, N, C]
o = np.matmul(beta, flatten_h) # [bs, N, C]
gamma = 0
o = np.reshape(o, X.shape) # [bs, h, w, C]
y = gamma * o + X
return y
But I have no idea how to add a trainable scalar gamma as described in the paper: SAGAN
I also hope someone can give some ideas about how to initialize a trainable keras scalar.
EDIT:
My implementation is now:
class Attention(Layer):
def __init__(self, ch, **kwargs):
super(Attention, self).__init__(**kwargs)
self.channels = ch
self.filters_f_g = self.channels // 8
self.filters_h = self.channels
def build(self, input_shape):
kernel_shape_f_g = (1, 1) + (self.channels, self.filters_f_g)
print(kernel_shape_f_g)
kernel_shape_h = (1, 1) + (self.channels, self.filters_h)
# Create a trainable weight variable for this layer:
self.gamma = self.add_weight(name='gamma', shape=[1], initializer='zeros', trainable=True)
self.kernel_f = self.add_weight(shape=kernel_shape_f_g,
initializer='glorot_uniform',
name='kernel_f')
self.kernel_g = self.add_weight(shape=kernel_shape_f_g,
initializer='glorot_uniform',
name='kernel_g')
self.kernel_h = self.add_weight(shape=kernel_shape_h,
initializer='glorot_uniform',
name='kernel_h')
self.bias_f = self.add_weight(shape=(self.filters_f_g,),
initializer='zeros',
name='bias_F')
self.bias_g = self.add_weight(shape=(self.filters_f_g,),
initializer='zeros',
name='bias_g')
self.bias_h = self.add_weight(shape=(self.filters_h,),
initializer='zeros',
name='bias_h')
super(Attention, self).build(input_shape)
# Set input spec.
self.input_spec = InputSpec(ndim=4,
axes={3: input_shape[-1]})
self.built = True
def call(self, x):
def hw_flatten(x):
return K.reshape(x, shape=[K.shape(x)[0], K.shape(x)[1]*K.shape(x)[2], K.shape(x)[-1]])
f = K.conv2d(x,
kernel=self.kernel_f,
strides=(1, 1), padding='same') # [bs, h, w, c']
f = K.bias_add(f, self.bias_f)
g = K.conv2d(x,
kernel=self.kernel_g,
strides=(1, 1), padding='same') # [bs, h, w, c']
g = K.bias_add(g, self.bias_g)
h = K.conv2d(x,
kernel=self.kernel_h,
strides=(1, 1), padding='same') # [bs, h, w, c]
h = K.bias_add(h, self.bias_h)
s = tf.matmul(hw_flatten(g), hw_flatten(f), transpose_b=True) # # [bs, N, N]
beta = K.softmax(s, axis=-1) # attention map
o = K.batch_dot(beta, hw_flatten(h)) # [bs, N, C]
o = K.reshape(o, shape=K.shape(x)) # [bs, h, w, C]
x = self.gamma * o + x
return x
def compute_output_shape(self, input_shape):
return input_shape

Answer 1

There are several problems with the modifications you made to the original code:
You cannot use numpy operations in the middle of your Keras/TF graph. First because numpy will try to operate directly, while the inputs tensors will actually be evaluated/receive their value only at graph runtime. Second because Keras/TF won't be able to back-propagate through non-Keras/TF operations.
You should replace the original tensorflow operations by their keras or keras.backend operations instead (e.g. tf.matmul() by keras.backend.batch_dot(), tf.nn.doftmax() by keras.backend.softmax(), etc.)
You are mixing Keras Layers (e.g. Conv2D) and Keras operations (e.g. np/keras.backend.reshape). Keras operations should be wrapped in a Lambda layer to be used along others.
Since this custom layer has a trainable parameter (gamma), you would need to write your own custom layer, e.g.:
from keras import backend as K
from keras.engine.topology import Layer
class AttentionLayer(Layer):
def __init__(self, output_dim, **kwargs):
self.output_dim = output_dim
super(AttentionLayer, self).__init__(**kwargs)
def build(self, input_shape):
# Create a trainable weight variable for this layer:
self.gamma = self.add_weight(name='gamma', shape=[1], initializer='uniform', trainable=True)
super(AttentionLayer, self).build(input_shape)
def call(self, x):
channels = K.int_shape(x)[-1]
x = activation(x, channels) # use TF implementation, or reimplement with Keras operations
return x
def compute_output_shape(self, input_shape):
return input_shape