tfe 模型保存和载入

原文链接: tfe 模型保存和载入

上一篇: tfe 配合 Keras model 线性拟合 和 自己处理梯度进行线性拟合

下一篇: tfe mnist 使用dataset 分类 保存和载入

简单参数保存和载入

如果路径不存在回自动创建

import tensorflow as tf
import tensorflow.contrib.eager as tfe

tf.enable_eager_execution()
x = tfe.Variable(10.)

checkpoint = tfe.Checkpoint(x=x)
x.assign(2.)  # Assign a new value to the variables and save.
print(x.numpy())  # 2.0

save_path = checkpoint.save('./ckpt/')
print(save_path)  # ./ckpt/-1

x.assign(11.)  # Change the variable after saving.
print(x.numpy())  # 11.0

# Restore values from the checkpoint
checkpoint.restore(save_path)
print(x.numpy())  # 2.0

使用Keras的Model时,需要保存很多参数,此时使用对象保存的方式

载入时使用的是模型的文件夹路径

主要代码

# 保存训练参数
optimizer = tf.train.AdamOptimizer(learning_rate=0.001)
checkpoint_dir = './save/'
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
root = tfe.Checkpoint(optimizer=optimizer,
                      model=model,
                      optimizer_step=tf.train.get_or_create_global_step())

root.save(file_prefix=checkpoint_prefix)
# or
# root.restore(tf.train.latest_checkpoint(checkpoint_dir))

完整代码

import tensorflow as tf
import tensorflow.contrib.eager as tfe
import os

tf.enable_eager_execution()


class Model(tf.keras.Model):
    def __init__(self):
        super(Model, self).__init__()
        self.W = tfe.Variable(5., name='weight')
        self.B = tfe.Variable(10., name='bias')

    def call(self, inputs):
        return inputs * self.W + self.B


# A toy dataset of points around 3 * x + 2
NUM_EXAMPLES = 2000
inputs = tf.random_normal([NUM_EXAMPLES])
noise = tf.random_normal([NUM_EXAMPLES])
targets = inputs * 3 + 2 + noise


# The loss function to be optimized
def loss():
    error = model(inputs) - targets
    return tf.reduce_mean(tf.square(error))


# Define:
# 1. A model.
# 2. Derivatives of a loss function with respect to model parameters.
# 3. A strategy for updating the variables based on the derivatives.
model = Model()
optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.01)

# 载入训练参数
# checkpoint_dir = './save/'
# checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
# root = tfe.Checkpoint(optimizer=optimizer,
#                       model=model,
#                       optimizer_step=tf.train.get_or_create_global_step())
# root.restore(tf.train.latest_checkpoint(checkpoint_dir))
#


# Training loop
for i in range(300):
    optimizer.minimize(loss)
    if i % 20 == 0:
        print(model.W.numpy(), model.B.numpy())

# 保存训练参数
optimizer = tf.train.AdamOptimizer(learning_rate=0.001)
checkpoint_dir = './save/'
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
root = tfe.Checkpoint(optimizer=optimizer,
                      model=model,
                      optimizer_step=tf.train.get_or_create_global_step())

root.save(file_prefix=checkpoint_prefix)
# or
# root.restore(tf.train.latest_checkpoint(checkpoint_dir))