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

在 parameters.py 中，定义了各类参数。

# training data directory
TRAINING_DATA_DIR = './data/'

# checkpoint directory
CHECKPOINT_DIR = './training_checkpoints/'

# training details
BATCH_SIZE = 16
BUFFER_SIZE = 128
EPOCHS = 15

---

在 numpy_dataset.py 中，创建了 5000 组训练数据集，模拟 y = x^3 + 1，并二进制格式写入文件。

from parameters import TRAINING_DATA_DIR

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


# create training data
X = np.linspace(-1, 1, 5000)
np.random.shuffle(X)
y = X ** 3 + 1 + np.random.normal(0, 0.01, (5000,))

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

# save data
if not os.path.exists(TRAINING_DATA_DIR):
    os.makedirs(TRAINING_DATA_DIR)

X.tofile(os.path.join(TRAINING_DATA_DIR + 'training_data_X.bin'))
y.tofile(os.path.join(TRAINING_DATA_DIR + 'training_data_y.bin'))

---

在 subclassed_model.py 中，通过对 tf.keras.models.Model 进行子类化，设计了两个自定义模型。

import tensorflow as tf
tf.enable_eager_execution()


# model definition
class Encoder(tf.keras.models.Model):
    def __init__(self):
        super(Encoder, self).__init__()
        self.fc1 = tf.keras.layers.Dense(units=16, activation='relu')
        self.fc2 = tf.keras.layers.Dense(units=8, activation='relu')

    def call(self, inputs):
        r = self.fc1(inputs)
        return self.fc2(r)


class Decoder(tf.keras.models.Model):
    def __init__(self):
        super(Decoder, self).__init__()
        self.fc = tf.keras.layers.Dense(units=1, activation=None)

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

---

在 loss_function.py 中，定义了损失函数。

import tensorflow as tf
tf.enable_eager_execution()


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

---

在 training.py 中，使用在 numpy_dataset.py 中创建的数据集训练模型，之后使用 model.save_weights() 保存 Keras Subclassed Model 模型，并创建验证集验证模型。

from parameters import TRAINING_DATA_DIR, CHECKPOINT_DIR, BATCH_SIZE, BUFFER_SIZE, EPOCHS
from subclassed_model import *
from loss_function import loss

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


# load training data
training_X = np.fromfile(os.path.join(TRAINING_DATA_DIR, 'training_data_X.bin'), dtype=np.float64)
training_y = np.fromfile(os.path.join(TRAINING_DATA_DIR, 'training_data_y.bin'), dtype=np.float64)

# plot training data
plt.scatter(training_X, training_y)
plt.show()

# training dataset
training_dataset = tf.data.Dataset.from_tensor_slices((training_X, training_y)).batch(BATCH_SIZE).shuffle(BUFFER_SIZE)

# model instance
encoder = Encoder()
decoder = Decoder()

# optimizer
optimizer = tf.train.AdamOptimizer()

# checkpoint
checkpoint_prefix_encoder = os.path.join(CHECKPOINT_DIR, 'encoder/', 'ckpt')
checkpoint_prefix_decoder = os.path.join(CHECKPOINT_DIR, 'decoder/', 'ckpt')

if not os.path.exists(os.path.dirname(checkpoint_prefix_encoder)):
    os.makedirs(os.path.dirname(checkpoint_prefix_encoder))
if not os.path.exists(os.path.dirname(checkpoint_prefix_decoder)):
    os.makedirs(os.path.dirname(checkpoint_prefix_decoder))

# training step
for epoch in range(EPOCHS):
    epoch_loss = 0

    for (batch, (tx, ty)) in enumerate(training_dataset):
        x = tf.cast(tx, tf.float32)
        y = tf.cast(ty, tf.float32)
        x = tf.expand_dims(x, axis=1)   # tf.Tensor([...], shape=(BATCH_SIZE, 1), dtype=float32)
        y = tf.expand_dims(y, axis=1)   # tf.Tensor([...], shape=(BATCH_SIZE, 1), dtype=float32)

        with tf.GradientTape() as tape:
            y_ = encoder(x)             # tf.Tensor([...], shape=(BATCH_SIZE, 8), dtype=float32)
            prediction = decoder(y_)    # tf.Tensor([...], shape=(BATCH_SIZE, 1), dtype=float32)
            batch_loss = loss(real=y, pred=prediction)

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

        epoch_loss += batch_loss

        if (batch + 1) % 100 == 0:
            print('Epoch {} Batch {} Loss {:.4f}'.format(epoch + 1,
                                                         batch + 1,
                                                         batch_loss.numpy()))

    print('Epoch {} Loss {:.4f}'.format(epoch + 1,
                                        epoch_loss / len(training_X)))

    if (epoch + 1) % 5 == 0:
        encoder.save_weights(checkpoint_prefix_encoder)
        decoder.save_weights(checkpoint_prefix_decoder)

# create evaluation data
X = np.linspace(-1, 1, 3000)
np.random.shuffle(X)

evaluation_X = tf.data.Dataset.from_tensor_slices(X).batch(BATCH_SIZE)
ey = []

for (batch, ex) in enumerate(evaluation_X):
    x = tf.cast(ex, tf.float32)
    x = tf.expand_dims(x, axis=1)
    prediction = decoder(encoder(x))
    for i in range(len(prediction.numpy())):
        ey.append(prediction.numpy()[i])

plt.scatter(X, ey)
plt.show()

# evaluate
eval_x = [[0.5]]
tensor_x = tf.convert_to_tensor(eval_x)
print(decoder(encoder(tensor_x)))

验证集评价结果如下图所示。

使用测试样例 eval_x 进行测试，测试结果如下。

tf.Tensor([[1.122567]], shape=(1, 1), dtype=float32)

---

在 evaluate.py 中，使用 model.load_weights() 恢复 Keras Subclassed Model 模型，并在验证集上进行验证，验证结果如下图所示。

from parameters import CHECKPOINT_DIR, BATCH_SIZE
from subclassed_model import *

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


# load model
enc = Encoder()
dec = Decoder()

enc.load_weights(tf.train.latest_checkpoint(os.path.join(CHECKPOINT_DIR, 'encoder/')))
dec.load_weights(tf.train.latest_checkpoint(os.path.join(CHECKPOINT_DIR, 'decoder/')))

# create evaluation data
X = np.linspace(-1, 1, 3000)
np.random.shuffle(X)

evaluation_X = tf.data.Dataset.from_tensor_slices(X).batch(BATCH_SIZE)
ey = []

for (batch, ex) in enumerate(evaluation_X):
    x = tf.cast(ex, tf.float32)
    x = tf.expand_dims(x, axis=1)
    prediction = dec(enc(x))
    for i in range(len(prediction.numpy())):
        ey.append(prediction.numpy()[i])

plt.scatter(X, ey)
plt.show()

# evaluate
eval_x = [[0.5]]
tensor_x = tf.convert_to_tensor(eval_x)
print(dec(enc(tensor_x)))

# model summary
enc.summary()
dec.summary()

使用测试样例 eval_x 进行测试，测试结果如下。

tf.Tensor([[1.122567]], shape=(1, 1), dtype=float32)

恢复模型的测试结果，与训练后模型的测试结果一致，且无需 build 模型。

---

版权声明：本文为博主原创文章，欢迎转载，转载请注明作者及原文出处！