Repository: GPflow/GPflow

Issue:
GPflow 2.0 & tf.saved_model
System information:

* Python version: 3.7
* TensorFlow installed from (source or binary): binary
* TensorFlow version (use command below): 'v2.0.0-rc2-26-g64c3d382ca', '2.0.0'
* GPflow installed from (source or binary): source
* GPflow version: develop_2.0, f070a8b

I tried to export a GPflow model using the `tf.saved_model` API by simply decorating the `predict_f` method with the `@tf.function` decorator. The export fails.

I'd expect the export to succeed.

The first obstacle seems to be the Parameter class. Even the following simple export fails:

```
import gpflow
import tensorflow as tf

x = tf.Variable(1.)
p = gpflow.base.Parameter(x)
tf.saved_model.save(p, "test")
```
Simply exporting the `tf.Variable` does work, exporting the `gpflow.base.Parameter` throws (Traceback below).

I've read the [tf.Module rfc](https://github.com/tensorflow/community/blob/master/rfcs/20190117-tf-module.md) and I gather that this is also used by GPflow [gpflow.base.Module](https://github.com/GPflow/GPflow/blob/develop-2.0/gpflow/base.py#L29) and [gpflow.base.Parameter](https://github.com/GPflow/GPflow/blob/develop-2.0/gpflow/base.py#L39). 

Is there any clear reasons why the export fails?

<pre>
Traceback (most recent call last):
  File "test.py", line 10, in <module>
    tf.saved_model.save(p, "test")
  File "C:\Users\madurno\AppData\Local\Continuum\anaconda3\envs\gpflow\lib\site-packages\tensorflow_core\python\saved_model\save.py", line 893, in save
    meta_graph_def, saveable_view, signatures, options.namespace_whitelist)
  File "C:\Users\madurno\AppData\Local\Continuum\anaconda3\envs\gpflow\lib\site-packages\tensorflow_core\python\saved_model\save.py", line 559, in _fill_meta_graph_def
    object_map, resource_map, asset_info = saveable_view.map_resources()
  File "C:\Users\madurno\AppData\Local\Continuum\anaconda3\envs\gpflow\lib\site-packages\tensorflow_core\python\saved_model\save.py", line 269, in map_resources
    new_variable = resource_variable_ops.copy_to_graph_uninitialized(obj)
  File "C:\Users\madurno\AppData\Local\Continuum\anaconda3\envs\gpflow\lib\site-packages\tensorflow_core\python\ops\resource_variable_ops.py", line 1949, in copy_to_graph_uninitialized
    constraint=var._constraint,
AttributeError: 'Parameter' object has no attribute '_constraint'
</pre>


Discussion:
Hello,

A bit of an update. The reason why it throws is because the Parameter class doesn't mirror all the properties of the Variable class. From what I can see the missing ones are: ` _constraint`, `shared_name`, `synchronization`, `aggregation` so the following patch is enough to make the code not throw:
```
--- a/gpflow/base.py
+++ b/gpflow/base.py
@@ -101,6 +101,22 @@ class Parameter(tf.Module):
     def trainable(self):
         return self._unconstrained.trainable

+    @property
+    def _constraint(self):
+        return self._unconstrained._constraint
+
+    @property
+    def _shared_name(self):
+        return self._unconstrained._shared_name
+
+    @property
+    def synchronization(self):
+        return self._unconstrained.synchronization
+
+    @property
+    def aggregation(self):
+        return self._unconstrained.aggregation
+
```
Saving the parameter no longer throws, however it is not possible to load the Parameter back.

### Saving a tf.Variable:
```
    x = tf.Variable([1.], dtype=tf.float64, name='variable')
    m = tf.Module()
    m.x = x
    tf.saved_model.save(m, "test_variable")
    m = tf.saved_model.load('./test_variable')
    print(m.x.read_value())
```
yields `tf.Tensor([1.], shape=(1,), dtype=float64)`

### Saving a gpflow.base.Parameter:
```
    x = tf.Variable([1.], dtype=tf.float64, name='variable')
    p = gpflow.base.Parameter(x, name='parameter')
    tf.saved_model.save(p, "test")
    p = tf.saved_model.load('./test')
    print(p._unconstrained.read_value())
```
successfully loads back the model but throws when reading the parameter value:
```
Traceback (most recent call last):
  File ".\test.py", line 19, in <module>
    print(p._unconstrained.read_value())
  File "C:\Users\madurno\AppData\Local\Continuum\anaconda3\envs\gpflow\lib\site-packages\tensorflow_core\python\ops\resource_variable_ops.py", line 630, in read_value
    value = self._read_variable_op()
  File "C:\Users\madurno\AppData\Local\Continuum\anaconda3\envs\gpflow\lib\site-packages\tensorflow_core\python\ops\resource_variable_ops.py", line 608, in _read_variable_op
    self._dtype)
  File "C:\Users\madurno\AppData\Local\Continuum\anaconda3\envs\gpflow\lib\site-packages\tensorflow_core\python\ops\gen_resource_variable_ops.py", line 582, in read_variable_op
    _six.raise_from(_core._status_to_exception(e.code, message), None)
  File "<string>", line 3, in raise_from
tensorflow.python.framework.errors_impl.FailedPreconditionError: Error while reading resource variable variable_218 from Container: localhost. This could mean that the variable was uninitialized. Not found: Resource localhost/variable_218/class tensorflow::Var does not exist. [Op:ReadVariableOp]
```

Is there any obvious reason why this happens?

Best Regards,

Marco
Hi Marco - sorry for the late reply here, thank you for opening the issue! We'll get on looking at this as soon as possible! I think the issue must have slipped under the radar...
Hello @marcoadurno, can you use checkpointing instead? https://www.tensorflow.org/api_docs/python/tf/train/Checkpoint
Hello @awav, yes checkpoint works fine but I need to be able to export the model as saved_model.

I have found a way around the tf bug by replacing the model parameters with variables holding the parameter value at the end of training. You lose the prior and the transform, but at least you can make predictions correctly reusing gpflow code and you can export the model as saved_model.

Having said that, I'm not suggesting we do something along those lines in gpflow as it's a bodge to workaround a TF bug.

Thanks,

Marco
Hi everyone, I am wondering if there are any news on this. I would also need to export models as saved models in my workflow. As I understand checkpointing only allows me to reload models in the same session.

@marcoadurno Could you let me know what exactly you do for your workaround? Are you manually replacing all parameters with variables, or is there a nicer solution?

Thanks in advance,
Max
Hi @mffrank ,

Yes, I simply replace all parameters with variables holding the parameter value once training is completed. You can do that in a programmatic way leveraging `gpflow.utilities.parameter_dict`.

Best Regards,

Marco
@mffrank this is now merged as `gpflow.utilities.freeze()` into develop (PR #1293).
Thanks!

Resolving patch:
diff --git a/doc/source/notebooks/intro_to_gpflow2.pct.py b/doc/source/notebooks/intro_to_gpflow2.pct.py
--- a/doc/source/notebooks/intro_to_gpflow2.pct.py
+++ b/doc/source/notebooks/intro_to_gpflow2.pct.py
@@ -6,7 +6,7 @@
 #       extension: .py
 #       format_name: percent
 #       format_version: '1.3'
-#       jupytext_version: 1.4.2
+#       jupytext_version: 1.6.0
 #   kernelspec:
 #     display_name: Python 3
 #     language: python
@@ -210,6 +210,8 @@ def noisy_sin(x):
 vgp_model.training_loss_closure(compile=False)  # uncompiled, same as vgp_model.training_loss
 
 # %% [markdown]
+# ### External data
+#
 # The SVGP model inherits from ExternalDataTrainingLossMixin and expects the data to be passed to training_loss().
 # For SVGP as for the other regression models, `data` is a two-tuple of `(X, Y)`, where `X` is an array/tensor with shape `(num_data, input_dim)` and `Y` is an array/tensor with shape `(num_data, output_dim)`:
 
@@ -473,7 +475,7 @@ def checkpointing_training_loop(
     )
 
 # %% [markdown]
-# ## Copying (hyper)parameter values between models
+# ### Copying (hyper)parameter values between models
 #
 # It is easy to interact with the set of all parameters of a model or a subcomponent programmatically.
 #
@@ -495,45 +497,32 @@ def checkpointing_training_loop(
 # %% [markdown]
 # ### TensorFlow `saved_model`
 #
-# At present, TensorFlow does not support saving custom variables like instances of the `gpflow.base.Parameter` class, see [this TensorFlow github issue](https://github.com/tensorflow/tensorflow/issues/34908).
-#
-# However, once training is complete, it is possible to clone the model and replace all `gpflow.base.Parameter`s with `tf.constant`s holding the same value:
-
-# %%
-model
-
-# %%
-frozen_model = gpflow.utilities.freeze(model)
-
-# %% [markdown]
-# In order to save the model we need to define a `tf.Module` holding the `tf.function`'s that we wish to export, as well as a reference to the underlying model:
+# In order to save the model we need to explicitly store the `tf.function`-compiled functions that we wish to export:
 
 # %%
-module_to_save = tf.Module()
-predict_fn = tf.function(
-    frozen_model.predict_f, input_signature=[tf.TensorSpec(shape=[None, 1], dtype=tf.float64)]
+model.predict_f_compiled = tf.function(
+    model.predict_f, input_signature=[tf.TensorSpec(shape=[None, 1], dtype=tf.float64)]
 )
-module_to_save.predict = predict_fn
 
 # %% [markdown]
-# Save original result for futher comparison. We also convert `samples_input` to a tensor. For a tensor input a `tf.function` will compile a single graph.
+# We also save the original prediction for later comparison. Here `samples_input` needs to be a tensor so that `tf.function` will compile a single graph:
 
 # %%
 samples_input = tf.convert_to_tensor(samples_input, dtype=default_float())
-original_result = module_to_save.predict(samples_input)
+original_result = model.predict_f_compiled(samples_input)
 
 # %% [markdown]
-# Let's save the module
+# Let's save the model:
 # %%
 save_dir = str(pathlib.Path(tempfile.gettempdir()))
-tf.saved_model.save(module_to_save, save_dir)
+tf.saved_model.save(model, save_dir)
 
 # %% [markdown]
-# Load module back as new instance and compare predict results
+# We can load the module back as a new instance and compare the prediction results:
 
 # %%
 loaded_model = tf.saved_model.load(save_dir)
-loaded_result = loaded_model.predict(samples_input)
+loaded_result = loaded_model.predict_f_compiled(samples_input)
 
 np.testing.assert_array_equal(loaded_result, original_result)
 


Test patch: