How to save tf.Variables for use as configuration parameters? - tensorflow

Question

So I’m working on a perceptual experiment for which I have designed a moderately complex stimulus which requires a good amount of computation to create. I had it implemented in Python; it works fine, but it’s slow to generate.
Quick brainstorm: hey, try implementing the stimulus computation in Tensorflow! A little bit of work paid off, cutting stimulus generation time in half (actually fast enough for real-time display), and the tensorflow model proves a more compact design. Great!
But then I started dreaming. I’ve wanted to be able to experiment via an app, but didn’t want to redo all my code in Swift or Java or whatever the language du jour for mobile/web is. But if all the heavy lifting code is embedded in a Tensorflow lite model, then a small wrapper for iOS/Android/(javascript maybe) would be doable in a reasonable time frame.
Here’s where my question comes in. In the off-line stimulus generation model, I configure my parameters and let it generate a video file, which can then be viewed by my subjects. If my theoretical app just takes a tensorflow model instead of a video file, then I’m really just shortening download time. What I’d really like to be able to do is to adjust the stimulus parameters within the app, instead of guessing, generating, and uploading again.
So (and from here on I’m just winging it as far as my TF skills) I turned my configuration parameters into tf.Variables, stuck them into the model, and voila, I can now adjust my stimulus on the fly, from within the Python CLI. Great! Now I just save the model…
Oops.
How does saving tf.Variables work? Here's a simple subset of my code to demonstrate the problem. Start with a layer that computes the sine of a temporal input, with a configurable phase, frequency, and amplitude:
class Sine(tf.keras.layers.Layer):
def __init__(self, *args, **kwargs):
super(Sine, self).__init__(*args, **kwargs)
self._twopi = tf.constant(np.pi * 2.0)
def call(self, parameters):
time = parameters[0]
scale = parameters[1]
frequency = parameters[2]
base = parameters[3]
phase = parameters[4]
# def call(self, time, scale, frequency, base, phase):
time = tf.cast(time, tf.float32)
return scale*tf.sin(self._twopi * frequency * time + phase) + base
Note here that I tried compressing the five arguments into one list to see what that would do.
Here's a stupid model that uses it:
class StupidModel:
def __init__(self, frequency, amplitude, base, phase):
self._frequency = tf.Variable(frequency, name="frequency", dtype=tf.float32)
self._amplitude = tf.Variable(amplitude, name="amplitude", dtype=tf.float32)
self._base = tf.Variable(base, name="base", dtype=tf.float32)
self._phase = tf.Variable(phase, name="phase", dtype=tf.float32)
self._model = self._build_model()
def _build_model(self):
input = tf.keras.layers.Input(1)
out = Sine()([input, self._frequency, self._amplitude, self._base, self._phase])
model = tf.keras.Model(inputs=[input], outputs = out)
model._myfrequency = self._frequency
model._myamplitude = self._amplitude
model._mybase = self._base
model._myphase = self._phase
return model
def __call__(self, time):
return self._model.predict(time)
And here's what happens when I try to save it:
>>> sm._model.save("foo.tf")
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
INFO:tensorflow:Assets written to: foo.tf/assets
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "/Users/j/.miniforge3/lib/python3.9/site-packages/keras/utils/traceback_utils.py", line 70, in error_handler
raise e.with_traceback(filtered_tb) from None
File "/Users/j/.miniforge3/lib/python3.9/json/encoder.py", line 199, in encode
chunks = self.iterencode(o, _one_shot=True)
File "/Users/j/.miniforge3/lib/python3.9/json/encoder.py", line 257, in iterencode
return _iterencode(o, 0)
TypeError: Unable to serialize <tf.Variable 'frequency:0' shape=() dtype=float32, numpy=0.5> to JSON. Unrecognized type <class 'tensorflow.python.ops.resource_variable_ops.ResourceVariable'>.
I believe that were I using TF1.X, I'd be looking at placeholders and feed dicts. But what can I do now? I know what I'm doing is probably outside of the normal tensorflow usage, e.g. no training, in-flight tuning, etc., but it works so well right up until saving...

Answer 1

Further exploration has yielded fruit. It seems complicated, and I shudder to think of what I'll need to do to convert to Tensorflow Lite, but as of now, I can:
Add arbitrary configuration parameters
Save the model
Load the model
Retrieve configuration parameters names
Set configuration parameters
The key I found was to make my StupidModel inherit from tf.Module. I added an intervening class to encapsulate the parameter logic:
class Parameterized(tf.Module):
def __init__(self, *args, **kwargs):
super(Parameterized, self).__init__()
def _save(self, name, variable):
self.__setattr__(name, variable)
def _initialize(self):
for key in self.variables().keys():
self.set(key, self.get(key))
def save(self, name):
tf.saved_model.save(self, name)
#tf.function
def variables(self):
return {k:self._trackable_children()[k] for k in self._trackable_children() if issubclass(type(self._trackable_children()[k]), tf.Variable)}
#tf.function
def set(self, variable, value):
self.__getattribute__(variable).assign(value)
#tf.function
def get(self, variable):
return self.__getattribute__(variable)
and then I updated the StupidModel to inherit from Parameterized, making sure to add the appropriate #tf.function decorator:
class StupidModel(Parameterized):
def __init__(self, frequency, amplitude, base, phase):
super().__init__(name="stupid")
self._save("frequency", tf.Variable(frequency, name="frequency", dtype=tf.float32))
self._save("amplitude", tf.Variable(amplitude, name="amplitude", dtype=tf.float32))
self._save("base", tf.Variable(base, name="base", dtype=tf.float32))
self._save("phase", tf.Variable(phase, name="phase", dtype=tf.float32))
self._sine = Sine()
self._initialize()
def _make_key(self, var):
return var.name.split(":")[0]
#tf.function(input_signature=(tf.TensorSpec(shape=[None], dtype=tf.float32),))
def __call__(self, time):
return self._sine(time, self.frequency, self.amplitude, self.base, self.phase)
Now, when I reload:
>>> import tensorflow as tf
>>> sm = tf.saved_model.load("foo.tf")
[...]
>>> sm([1, 2, 3])
<tf.Tensor: shape=(3,), dtype=float32, numpy=array([0.2602872 , 0.73971283, 0.26028723], dtype=float32)>
And I can check out my parameters (by name), set them, and see the change:
>>> sm.variables().keys()
dict_keys(['phase', 'amplitude', 'base', 'frequency'])
>>> sm.set('base', 5.0)
>>> sm([1, 2, 3])
<tf.Tensor: shape=(3,), dtype=float32, numpy=array([4.7602873, 5.2397127, 4.7602873], dtype=float32)>
And one final important note: the tf.function argument provided allows me to use arbitrary inputs. (Without it, I can only call __call__ with an argument signature that exactly matches calls made before the save, e.g. if I called sm([1, 2, 3]) and then saved, I'd get an error back if I called sm([10, 20, 30]) after reloading. Specifying the signature in #tf.function prevents that.):
>>> sm([10, 20, 30])
<tf.Tensor: shape=(3,), dtype=float32, numpy=array([5.239713 , 5.239714 , 5.2397127], dtype=float32)>
So that's what I've learned. It seems complicated, and I'm definitely open to easier approaches.