tensorflow checkpoint设置检查点

https://tensorflow.google.cn/guide/checkpoint

import tensorflow as tf

# 模型
class Net(tf.keras.Model):
    """A simple linear model."""

    def __init__(self):
        super(Net, self).__init__()
        self.l1 = tf.keras.layers.Dense(5)
    def call(self, x):
        return self.l1(x)
net = Net()
# keras保存权重
net.save_weights('easy_checkpoint')  # 从 tf.keras 训练 API 保存

# 加载数据
def toy_dataset():
    inputs = tf.range(10.)[:, None]
    labels = inputs * 5. + tf.range(5.)[None, :]
    return tf.data.Dataset.from_tensor_slices(
        dict(x=inputs, y=labels)).repeat().batch(2)
dataset = toy_dataset()

# 更新梯度步骤
def train_step(net, example, optimizer):
    """Trains `net` on `example` using `optimizer`."""
    with tf.GradientTape() as tape:
        output = net(example['x'])
        loss = tf.reduce_mean(tf.abs(output - example['y']))
    variables = net.trainable_variables
    gradients = tape.gradient(loss, variables)
    optimizer.apply_gradients(zip(gradients, variables))
    return loss

# 优化器
opt = tf.keras.optimizers.Adam(0.1)
# 迭代数据
iterator = iter(dataset)
# 设置检查点
ckpt = tf.train.Checkpoint(step=tf.Variable(1), optimizer=opt, net=net, iterator=iterator)
# 设置检查点管理
manager = tf.train.CheckpointManager(ckpt, './tf_ckpts', max_to_keep=3)

# 输入训练网络和检查点管理，开始训练
def train_and_checkpoint(net, manager):
    ckpt.restore(manager.latest_checkpoint)
    if manager.latest_checkpoint:
        print("Restored from {}".format(manager.latest_checkpoint))
    else:
        print("Initializing from scratch.")

    for _ in range(50):
        example = next(iterator)
        loss = train_step(net, example, opt)
        ckpt.step.assign_add(1)
        if int(ckpt.step) % 10 == 0:
            save_path = manager.save()
            print("Saved checkpoint for step {}: {}".format(int(ckpt.step), save_path))
            print("loss {:1.2f}".format(loss.numpy()))
train_and_checkpoint(net, manager)

https://tensorflow.google.cn/guide/keras/save_and_serialize

如果使用keras建模的可以使用 tf.keras 保存和恢复指南。

https://tensorflow.google.cn/guide/keras/save_and_serialize

tf.keras.Model.save_weights 可以保存一个 TensorFlow 检查点。

10秒了解

# 定义模型 Get model (Sequential, Functional Model, or Model subclass)
model = ...
# 保存模型  
model.save('path/to/location')
# 加载模型
from tensorflow import keras
model = keras.models.load_model('path/to/location')

save了哪些

• 模型的架构/配置
• 模型的权重值（在训练过程中学习）
• 模型的编译信息（如果调用了 compile()）
• 优化器及其状态（如果有的话，使您可以从上次中断的位置重新开始训练）

水平有限，貌似GAN如果用的keras架构的不能直接save-load，还是要用检查点。参考下面两个教程。

https://tensorflow.google.cn/tutorials/generative/pix2pix

Deep Convolutional Generative Adversarial Network  |  TensorFlow Core