在 Tensorflow 2 中为 Model(object) 类创建检查点

Question

I was just trying to use low level API of tensorflow2. I created my model based on this tutorial : https://www.tensorflow.org/tutorials/customization/custom_training#define_the_model

Then I want to create a checkpoint for my training process, and I follow this tutorial : https://www.tensorflow.org/guide/checkpoint

The problem is the checkpoint's tutorial use a class with tf.keras.Model as parameter, while I use object as my parameter. It gave me error, said that it was expecting a trackable object.

Here is the snippet of my code:

class SimpleANN(object):
    def __init__(self):
        initializer = tf.initializers.glorot_uniform()
        self.w1 = tf.Variable(initializer([784, 360]), name = 'weight1', trainable = True, dtype = tf.float32)
        self.w2 = tf.Variable(initializer([360, 64]), name = 'weight2', trainable = True, dtype = tf.float32)
        self.w3 = tf.Variable(initializer([64, 10]), name = 'weight3', trainable = True, dtype = tf.float32)

    def __call__(self, x, leaky_relu_alpha = 0.2):
        fc1 = tf.nn.leaky_relu(tf.matmul(x, self.w1), alpha = leaky_relu_alpha)
        fc2 = tf.nn.leaky_relu(tf.matmul(fc1, self.w2), alpha = leaky_relu_alpha)
        logits = tf.matmul(fc2, self.w3)

        return logits

model = SimpleANN() 
optimizer = tf.keras.optimizers.Adam(learning_rate)

ckpt = tf.train.Checkpoint(step=tf.Variable(1), optimizer=optimizer, model = model)

then I got this error :

ValueError: `Checkpoint` was expecting a trackable object (an object
derived from `TrackableBase`), got <__main__.SimpleANN object at
0x000001D859792748>. 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.

I would like to know If it is able to implement tf.train.Checkpoint for the low level API, as what I was doing.

Answer 1

What you try is far less, then a tensorflow keras Model does. If you really want to know, what does "low level" model creation mean, take a visit to github: tensorflow/tensorflow/python/keras/engine/training.py , where you will find a lot reference to of what you need to do to reach success with your approach eg @trackable.no_automatic_dependency_tracking decorator. As you will see there, even Model() class has parameters, you have to dive into. Obviously it is not impossible, but you have to dive deeper.

Answer 2

The TensorFlow SavedModel API can only be used to save trackable objects (which Keras models are by default). One way to create a trackable object using the low-level API is to inherit from tf.Module . In my environment (Python version 3.7.6, TensorFlow version 2.1.0), I can make the errors go away by replacing the line class SimpleANN(object): with class SimpleANN(tf.Module): .

You may or may not also need to decorate your methods with @tf.function , and call them once to trace-compile a graph before saving a checkpoint. For more information, see my answer to RobR's question here .