Custom optimiser implementation in Keras - tensorflow

Question

I am trying to use a custom optimiser to train a NN in Keras. The original algorithm has been developed to train CNNs and is based on using different adaptive learning rates for every weight of the network. The algorithm is called WAME (weight-wise adaptive learning rates with moving average estimator).
The optimiser has been developed by a former University of London student and can be found in this GitHub repo (lines 54 to 153).
As you can see in the code, it is built as a subclass of the optimizer_v2 superclass available in TensorFlow. This class is called WAMEprop.
What I am trying to do is simply:
pasting the code that defines the class into my Google Colab notebook
using the new optimizer as follows:
#building the model
wame_model = models.Sequential()
wame_model.add(layers.Dense(32, activation='relu', input_shape=(11,)))
wame_model.add(layers.Dense(32, activation='relu'))
wame_model.add(layers.Dense(1))
#compiling the model using WAMEprop as optimizer
wame_model.compile(optimizer=WAMEprop(name='wame'), loss='mse', metrics=['mae'])
#fitting the model
history = wame_model.fit(train_features, train_targets,
validation_data=(test_features, test_targets),
epochs=50, batch_size=1, verbose=1)
Now, I get the following error:
ValueError Traceback (most recent call last)
<ipython-input-55-a13c54fc61f2> in <module>()
5 wame_model.add(layers.Dense(32, activation='relu'))
6 wame_model.add(layers.Dense(1))
----> 7 wame_model.compile(optimizer=WAMEprop(name='wame'), loss='mse', metrics=['mae'])
8
9 history = wame_model.fit(white_train_features_df_s.to_numpy(), white_train_targets.to_numpy(),
5 frames
/usr/local/lib/python3.7/dist-packages/keras/optimizers.py in get(identifier)
131 else:
132 raise ValueError(
--> 133 'Could not interpret optimizer identifier: {}'.format(identifier))
ValueError: Could not interpret optimizer identifier: <__main__.WAMEprop object at 0x7fd4bc1d01d0>
Since I am new to Keras and I, unfortunately, don't know Tensorflow, I am not sure what exactly it is not able to find.
Did I get some import wrong?
Did I use the keras compile() method in the wrong way?
Also, if I don't pass the parameter name='wame' to the WAMEprop() call, I get an error message saying that the positional argument 'name' is required. Strangely enough, there is no parameter 'name' in the class constructor. Does this depend on the interaction with the superclass?
Thank you very much a lot in advance for any help you could offer!
Cheers!
UPDATE:
the error message refers to a method get() that takes as input identifier in the optimizers.py file that must have been installed with TensorFlow. Now, this function is expecting to get a string (I guess fr the readily available optimizers), a configuration dictionary, an optimizer_v2.OptimizerV2 object or a tf.compat.v1.train.Optimizer object.
I think the object I am passing as an optimizer is no one of these.
If I run:
my_optimizer = WAMEprop(name='wame')
print(type(my_optimizer))
I get <class '__main__.WAMEprop'>.
So, I suspect the object I am dealing with is something different from what Keras is expecting.
UPDATE2: it runs on my laptop, I have tensorflow installed within an Anaconda environment. Now, I am convinced there is some installation or import problem in Google Colab

Answer 1

Problem solved. I needed to uninstall tensorflow 2.6.0 from Google Colab and install tensorflow 2.0.0 instead.