Tensorflow 2.2.0 and Keras save model / load model problems - tensorflow

Question

After adding a custom loss function as #tf.function to my keras DQN, keras models stopped loading (seems to save model, but cannot reload model). Documentation suggests this is really simple, but...
Various SO answers suggest that models trained using one Keras version cannot be loaded into other Keras version. So I uninstalled Keras 2.4.3 (from Anaconda env), to avoid any confusion, and trying to solely model and save/load using Tensorflow-keras.
So, now trying to save a Tensorflow-keras model and then load that model again, but will not re-load, various errors (below). Environment is Anaconda3 python3.8 (with Keras 2.4.3, then uninstalled this) and Tensorflow 2.2.0 (containing Keras 2.3.0-tf).
Is there some solution to simply save a model and then reload a model in tf 2.2.0 (with keras 2.3.0-tf)?
import tensorflow as tf
from tensorflow.keras.models import Sequential, load_model
from tensorflow.keras.layers import Dense, LSTM, Masking, Input
from tensorflow.keras.optimizers import Adam
Then all tf.keras modelling and save / load should be done by Keras-2.3.0-tf, from within Tensorflow. Model save is done with:
agent.model.save(os.path.join(pathOUT, PAIR, 'models' + modelNum, modelFolder),
save_format='tf')
But generates deprecation warning during save:
2020-11-26 00:19:03.388858: W tensorflow/python/util/util.cc:329] Sets are not currently
considered sequences, but this may change in the future, so consider avoiding using them.
WARNING:tensorflow:From C:\..mypath......\lib\site-
packages\tensorflow\python\ops\resource_variable_ops.py:1813: calling
BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with
constraint is deprecated and will be removed in a future version.
Instructions for updating:
If using Keras pass *_constraint arguments to layers.
Saved an Intermediate model...
Then attempt to load model with:
model = load_model(LOAD_MODEL)
but generates error during loading:
TypeError: __init__() got an unexpected keyword argument 'reduction'
Again, is there some solution to simply save a model and then reload a model in tf 2.2.0 (with keras 2.3.0-tf)?
Full error:
Traceback (most recent call last):
File mypath, line 851, in <module>
agent = DQNAgent()
File mypath, line 266, in __init__
self.model = self.create_model()
File mypath, line 336, in create_model
model = load_model(LOAD_MODEL)
File mypath\lib\site-packages\tensorflow\python\keras\saving\save.py", line 190, in load_model
return saved_model_load.load(filepath, compile)
File mypath\lib\site-packages\tensorflow\python\keras\saving\saved_model\load.py", line 116,
in load
model = tf_load.load_internal(path, loader_cls=KerasObjectLoader)
File mypath\lib\site-packages\tensorflow\python\saved_model\load.py", line 602, in
load_internal
loader = loader_cls(object_graph_proto,
File mypath\lib\site-packages\tensorflow\python\keras\saving\saved_model\load.py", line 188,
in __init__
super(KerasObjectLoader, self).__init__(*args, **kwargs)
File mypath\lib\site-packages\tensorflow\python\saved_model\load.py", line 123, in __init__
self._load_all()
File mypath\lib\site-packages\tensorflow\python\keras\saving\saved_model\load.py", line 209,
in _load_all
self._layer_nodes = self._load_layers()
File mypath\lib\site-packages\tensorflow\python\keras\saving\saved_model\load.py", line 312,
in _load_layers
layers[node_id] = self._load_layer(proto.user_object, node_id)
File mypath\lib\site-packages\tensorflow\python\keras\saving\saved_model\load.py", line 335,
in _load_layer
obj, setter = self._revive_from_config(proto.identifier, metadata, node_id)
File mypath\lib\site-packages\tensorflow\python\keras\saving\saved_model\load.py", line 349,
in _revive_from_config
obj = self._revive_metric_from_config(metadata, node_id)
File mypath\lib\site-packages\tensorflow\python\keras\saving\saved_model\load.py", line 441,
in _revive_metric_from_config
obj = metrics.deserialize(
File mypath\lib\site-packages\tensorflow\python\keras\metrics.py", line 3345, in deserialize
return deserialize_keras_object(
File mypath\lib\site-packages\tensorflow\python\keras\utils\generic_utils.py", line 361, in
deserialize_keras_object
(cls, cls_config) = class_and_config_for_serialized_keras_object(
File mypath\lib\site-packages\tensorflow\python\keras\utils\generic_utils.py", line 327, in
class_and_config_for_serialized_keras_object
deserialized_objects[key] = deserialize_keras_object(
File mypath\lib\site-packages\tensorflow\python\keras\utils\generic_utils.py", line 375, in
deserialize_keras_object
return cls.from_config(cls_config)
File mypath\lib\site-packages\tensorflow\python\keras\metrics.py", line 628, in from_config
return super(MeanMetricWrapper, cls).from_config(config)
File mypath\lib\site-packages\tensorflow\python\keras\engine\base_layer.py", line 655, in
from_config
return cls(**config)
TypeError: __init__() got an unexpected keyword argument 'reduction'
Additional code:
Custom loss function (trying to implement gradient ascent, by 'flipping' the error gradient), have tried this in various locations (within same agent class as model, outside agent class, other.
#tf.function
def positive_mse(y_true, y_pred):
return -1 * tf.keras.losses.MSE(y_true, y_pred)
The Sequential model I am using is here

Answer 1

Have resolved/bypassed the original keyword argument 'reduction' error by REMOVING the MeanSquareError() metric during model compile of original model. Original model:
model.compile(loss=positive_mse,
optimizer=Adam(lr=LEARNING_RATE, decay=DECAY),
metrics=[tf.keras.losses.MeanSquaredError()])
From the Keras docs: "Note that this is an important difference between loss functions like tf.keras.losses.mean_squared_error and default loss class instances like tf.keras.losses.MeanSquaredError: the function version does not perform reduction, but by default the class instance does."
The MeanSquaaredError loss class function is passing a 'reduction' keyword during evaluation of loss over a minibatch. Removing this metric allows model to be reloaded without error.