[Tensorflow] 使用 model.save_weights() 保存 Keras Subclassed Model

 import numpy as np
 import matplotlib.pyplot as plt
 import os
 import time

 import tensorflow as tf
 tf.enable_eager_execution()

 # create data
 X = np.linspace(-1, 1, 5000)
 np.random.shuffle(X)
 y = 0.5 * X + 2 + np.random.normal(0, 0.05, (5000,))

 # plot data
 plt.scatter(X, y)
 plt.show()

 # split data
 X_train, y_train = X[:4000], y[:4000]
 X_test, y_test = X[4000:], y[4000:]

 # tf.data
 BATCH_SIZE = 32
 BUFFER_SIZE = 512
 dataset = tf.data.Dataset.from_tensor_slices((X_train, y_train)).batch(BATCH_SIZE).shuffle(BUFFER_SIZE)

 # subclassed model
 UNITS = 1

 class Model(tf.keras.Model):
     def __init__(self):
         super(Model, self).__init__()
         self.fc = tf.keras.layers.Dense(units=UNITS)

     def call(self, inputs):
         return self.fc(inputs)

 model = Model()

 optimizer = tf.train.AdamOptimizer()

 # loss function
 def loss_function(real, pred):
     return tf.losses.mean_squared_error(labels=real, predictions=pred)

 EPOCHS = 30
 checkpoint_dir = './save_subclassed_keras_model_training_checkpoints'
 if not os.path.exists(checkpoint_dir):
     os.makedirs(checkpoint_dir)

 # training loop
 for epoch in range(EPOCHS):
     start = time.time()
     epoch_loss = 0

     for (batch, (x, y)) in enumerate(dataset):
         x = tf.cast(x, tf.float32)
         y = tf.cast(y, tf.float32)
         x = tf.expand_dims(x, axis=1)
         y = tf.expand_dims(y, axis=1)
         # print(x)    # tf.Tensor([...], shape=(BATCH_SIZE, 1), dtype=float32)
         # print(y)    # tf.Tensor([...], shape=(BATCH_SIZE, 1), dtype=float32)
         with tf.GradientTape() as tape:
             predictions = model(x)
             # print(predictions)  # tf.Tensor([...], shape=(BATCH_SIZE, 1), dtype=float32)
             batch_loss = loss_function(real=y, pred=predictions)

         grads = tape.gradient(batch_loss, model.variables)
         optimizer.apply_gradients(zip(grads, model.variables),
                                   global_step=tf.train.get_or_create_global_step())
         epoch_loss += batch_loss

         if (batch + 1) % 10 == 0:
             print('Epoch {} Batch {} Loss {:.4f}'.format(epoch + 1,
                                                          batch + 1,
                                                          batch_loss/int(x.shape[0])))

     print('Epoch {} Loss {:.4f}'.format(epoch + 1, epoch_loss/len(X_train)))
     print('Time taken for 1 epoch {} sec\n'.format(time.time() - start))

     # save checkpoint
     checkpoint_prefix = os.path.join(checkpoint_dir, 'ckpt')
     if (epoch + 1) % 10 == 0:
         model.save_weights(checkpoint_prefix.format(epoch=epoch), overwrite=True)

 _model = Model()
 _model.load_weights(tf.train.latest_checkpoint(checkpoint_dir))
 _model.build(input_shape=tf.TensorShape([BATCH_SIZE, 1]))
 _model.summary()

 test_dataset = tf.data.Dataset.from_tensor_slices(X_test).batch(1)

 for (batch, x) in enumerate(test_dataset):
     x = tf.cast(x, tf.float32)
     x = tf.expand_dims(x, axis=1)
     print(x)
     predictions = _model(x)
     print(predictions)
     exit()