How to save variable in custom layer in keras - tensorflow

Question

I wI want to create a (4,1) dimension matrix to multiply with inputs with (4,) dimension in my custom layer. So I do as the following:
class MyLayer(Layer):
def __init__(self, output_dim = 4, **kwargs):
self.output_dim = output_dim
super(MyLayer, self).__init__(**kwargs)
def build(self, input_shape):
# Create a trainable weight variable for this layer.
k1 = K.variable(1, dtype=tf.float32, name='k1')
k2 = K.variable(1, dtype=tf.float32, name='k2')
k3 = K.variable(1, dtype=tf.float32, name='k3')
k4 = K.variable(1, dtype=tf.float32, name='k4')
weights = tf.stack([k1,k2,k3,k4], name='weight_stack')
weights = tf.expand_dims(weights,1)
weights = tf.expand_dims(weights, 1)
weights = tf.expand_dims(weights, 1)
weights = tf.expand_dims(weights, 1)
self.kernel=tf.nn.softmax(weights,name='weights_softmax')
super(MyLayer, self).build(input_shape) # Be sure to call this somewhere!
def call(self, inputs):
self.inputs = inputs
P = tf.multiply(self.inputs, self.kernel)
P = tf.add_n([P[0],P[1],P[2],P[3]])
self.shape = P.shape
return P
def compute_output_shape(self, input_shape ):
return (input_shape, 1, 1, 1, 1)
I call it as followed:
x = [x1, x2, x3, x4]
x = MyLayer(name = "weight_matrix")(x)
It works when I train the Network with this self designed layers. However, when I save it with callbacks as followed:
keras.callbacks.ModelCheckpoint(self.checkpoint_path,
verbose=0, save_weights_only=True)
I found that I can't reload the "weights"(means k1 to k4) when I try to reload the model.
And I take a look into the file which save the model, I can't find the weights for the "weight_matrix".
If I am not doing it right, how can I fix it or how to implement it in keras? Thx.

Answer 1

Nowhere you're calling self.add_weight.
Your layer simply doesn't have weights.
In build:
def build(self, input_shape):
def customWeights(shape, dtype=None):
#k1 = K.variable(1, dtype=tf.float32, name='k1')
#k2 = K.variable(1, dtype=tf.float32, name='k2')
#k3 = K.variable(1, dtype=tf.float32, name='k3')
#k4 = K.variable(1, dtype=tf.float32, name='k4')
#weights = tf.stack([k1,k2,k3,k4], name='weight_stack')
weights = K.ones((4,))
#you can also use the shape passed to this function to build a flexible layer
#weights = tf.expand_dims(weights,1)
#weights = tf.expand_dims(weights, 1)
#weights = tf.expand_dims(weights, 1)
#weights = tf.expand_dims(weights, 1)
weights = K.reshape(weights,(-1,1,1,1,1))
#this should be applied later
#return tf.nn.softmax(weights,name='weights_softmax')
return weights
self.kernel= self.add_weight(name='kernel',
shape=notBeingUsedButItMayBeGoodForFlexibleLayers,
initializer=customWeights,
trainable=True)
super(MyLayer, self).build(input_shape) # Be sure to call this somewhere!
In call you call the softmax on your weights:
def call(self, inputs):
#don't call this! You're replacing an important property of the layer
#self.inputs = inputs
kernel = K.softmax(self.kernel)
P = tf.multiply(inputs, kernel)
P = tf.add_n([P[0],P[1],P[2],P[3]])
#not sure if this is a good idea either
#self.shape = P.shape
return P