tf.train.Checkpoint does not support RMSprop weights for tensorflow-macos 2.8.8

[MarkTension]

Click to expand!

Issue Type

Feature Request

Source

source

Tensorflow Version

tensorflow-macos 2.8.0

Custom Code

No

OS Platform and Distribution

MacOS

Mobile device

No response

Python version

3.8.8

Bazel version

No response

GCC/Compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current Behaviour?

Optimizer is not supported

When I select 'RMSProp' in 
checkpoint = tf.train.Checkpoint(optimizer='RMSprop', model=model)
Then I get (1). The same when I don't use a string, but a full optimizer declaration in keras 

When I run it without setting an optimizer, then status.assert_consumed() starts complaining:
(2)

Standalone code to reproduce the issue

keras.optimizer_v2.rmsprop.RMSprop 

tf.train.Checkpoint(optimizer='RMSprop', model=model)

Relevant log output

(1)
Traceback (most recent call last):
  File "main.py", line 80, in <module>
    main(config_dot)
  File "main.py", line 47, in main
    model = load_model(config, training_data)
  File "/Users/mark/Code/pythonProjects/music_RNN/train.py", line 127, in load_model
    checkpoint = tf.train.Checkpoint(optimizer='RMSprop', model=model)
  File "/Users/mark/opt/anaconda3/envs/tensorlfowGPU/lib/python3.8/site-packages/tensorflow/python/training/tracking/util.py", line 2017, in __init__
    _assert_trackable(converted_v, k)
  File "/Users/mark/opt/anaconda3/envs/tensorlfowGPU/lib/python3.8/site-packages/tensorflow/python/training/tracking/util.py", line 1463, in _assert_trackable
    raise ValueError(
ValueError: `Checkpoint` was expecting optimizer to be a trackable object (an object derived from `Trackable`), got RMSprop. If you believe this object should be trackable (i.e. it is part of the TensorFlow Python API and manages state), please open an issue.



(2)
Traceback (most recent call last):
  File "main.py", line 80, in <module>
    main(config_dot)
  File "main.py", line 47, in main
    model = load_model(config, training_data)
  File "/Users/mark/Code/pythonProjects/music_RNN/train.py", line 131, in load_model
    status.assert_consumed()
  File "/Users/mark/opt/anaconda3/envs/tensorlfowGPU/lib/python3.8/site-packages/tensorflow/python/training/tracking/util.py", line 784, in assert_consumed
    raise AssertionError(
AssertionError: Unresolved object in checkpoint (root).model.optimizer.iter: attributes {
  name: "VARIABLE_VALUE"
  full_name: "RMSprop/iter"
  checkpoint_key: "model/optimizer/iter/.ATTRIBUTES/VARIABLE_VALUE"
}
has_checkpoint_values {
  value: true
}

WARNING:tensorflow:Detecting that an object or model or tf.train.Checkpoint is being deleted with unrestored values. See the following logs for the specific values in question. To silence these warnings, use `status.expect_partial()`. See https://www.tensorflow.org/api_docs/python/tf/train/Checkpoint#restorefor details about the status object returned by the restore function.
WARNING:tensorflow:Value in checkpoint could not be found in the restored object: (root).model.optimizer.iter
WARNING:tensorflow:Value in checkpoint could not be found in the restored object: (root).model.optimizer.decay
WARNING:tensorflow:Value in checkpoint could not be found in the restored object: (root).model.optimizer.learning_rate
WARNING:tensorflow:Value in checkpoint could not be found in the restored object: (root).model.optimizer.momentum
WARNING:tensorflow:Value in checkpoint could not be found in the restored object: (root).model.optimizer.rho

[gowthamkpr]

@MarkTension This is not specific to MacOS. The tf.train.Checkpoint() receives kwargs paras which are key-object. And the object must be a trackable object(RMSprop is not a trackable object)

If you replace optimizer RMSprop with tf.keras.optimizers.RMSprop(0.1), it works. Heres the code.

[google-ml-butler]

This issue has been automatically marked as stale because it has no recent activity. It will be closed if no further activity occurs. Thank you.

[google-ml-butler]

Closing as stale. Please reopen if you'd like to work on this further.