TensorFlow自定义模型保存加载和分布式训练

class CustomModel(tf.keras.Model):
    def __init__(self):
        super(CustomModel, self).__init__()
        self.layer1 = tf.keras.layers.Dense(5, activation='relu')
        self.layer2 = tf.keras.layers.Dense(1, activation='sigmoid')
    def call(self, inputs):
        x = self.layer1(inputs)
        return self.layer2(x)
model = CustomModel()
# 定义优化器和损失函数
optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)
loss_fn = tf.keras.losses.BinaryCrossentropy()
# 创建 Checkpoint
ckpt = tf.train.Checkpoint(step=tf.Variable(1), optimizer=optimizer, model=model)
# 训练模型
# ...
# 保存模型
ckpt.save('/path/to/ckpt')
# 加载模型
ckpt.restore(tf.train.latest_checkpoint('/path/to/ckpt'))

# 创建一个 MirroredStrategy 对象
strategy = tf.distribute.MirroredStrategy()

with strategy.scope():
    # 在策略范围内创建模型和优化器
    model = CustomModel()
    optimizer = tf.keras.optimizers.Adam()

    loss_fn = tf.keras.losses.BinaryCrossentropy()
    metrics = [tf.keras.metrics.Accuracy()]

    model.compile(optimizer=optimizer, loss=loss_fn, metrics=metrics)

# 在所有可用的设备上训练模型
model.fit(train_dataset, epochs=10)

import tensorflow as tf
from tensorflow.keras.callbacks import TensorBoard
# 创建一个简单的模型
model = tf.keras.models.Sequential([
    tf.keras.layers.Dense(32, activation='relu', input_shape=(100,)),
    tf.keras.layers.Dense(1, activation='sigmoid')
])
# 编译模型
model.compile(optimizer='adam',
              loss='binary_crossentropy',
              metrics=['accuracy'])
# 创建一个 TensorBoard 回调
tensorboard_callback = TensorBoard(log_dir='./logs', histogram_freq=1)
# 使用训练数据集训练模型，并通过验证数据集验证模型
model.fit(train_dataset,
          epochs=5,
          validation_data=validation_dataset,
          callbacks=[tensorboard_callback])