带有tf.keras的Hparams插件(TensorFlow2.0)

import tensorflow as tf
from tensorboard.plugins.hparams import api as hp

# 定义超参数空间
HP_NUM_UNITS = hp.HParam('num_units', hp.Discrete([16, 32, 64]))
HP_LEARNING_RATE = hp.HParam('learning_rate', hp.RealInterval(0.001, 0.01))

# 构建模型
def build_model(hparams):
    model = tf.keras.Sequential([
        tf.keras.layers.Dense(hparams[HP_NUM_UNITS], activation='relu'),
        tf.keras.layers.Dense(10, activation='softmax')
    ])
    optimizer = tf.keras.optimizers.Adam(hparams[HP_LEARNING_RATE])
    model.compile(optimizer=optimizer, loss='sparse_categorical_crossentropy', metrics=['accuracy'])
    return model

# 训练模型
def train_model(hparams, train_data, train_labels):
    model = build_model(hparams)
    model.fit(train_data, train_labels, epochs=5)
    return model

# 定义实验
with tf.summary.create_file_writer('logs/hparam_tuning').as_default():
    hp.hparams_config(
        hparams=[HP_NUM_UNITS, HP_LEARNING_RATE],
        metrics=[hp.Metric('accuracy', display_name='Accuracy')]
    )

# 运行实验
session_num = 0
for num_units in HP_NUM_UNITS.domain.values:
    for learning_rate in (HP_LEARNING_RATE.domain.min_value, HP_LEARNING_RATE.domain.max_value):
        hparams = {
            HP_NUM_UNITS: num_units,
            HP_LEARNING_RATE: learning_rate
        }
        model = train_model(hparams, train_data, train_labels)
        accuracy = model.evaluate(test_data, test_labels)[1]
        tf.summary.scalar('accuracy', accuracy, step=session_num)
        hp.hparams(hparams, step=session_num)
        session_num += 1