原文链接: 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))