我正在尝试在 TensorFlow 2.0 alpha 上进行自定义训练，同时我正在尝试将一些指标和我的训练图添加到 TensorBoard。考虑以下人为的示例
import tensorflow as tf
from tensorflow.keras.layers import Dense, Input
from tensorflow.keras.models import Model
def create_model():
inp = Input((32, ))
net = Dense(16, activation="relu")(inp)
net = Dense(8, activation="relu")(net)
net = Dense(2, activation=None)(net)
return Model(inp, net)
@tf.function
def grad(model, loss, x, y):
with tf.GradientTape() as tape:
y_ = model(x)
loss_value = loss(y_true=y, y_pred=y_)
return loss_value, tape.gradient(loss_value, model.trainable_variables)
@tf.function
def train_step(model, loss, optimizer, features, labels):
loss_value, grads = grad(model, loss, features, labels)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
return loss_value
def train():
tf.summary.trace_on(graph=True, profiler=True)
with tf.summary.create_file_writer("model").as_default():
model = create_model()
loss = tf.keras.losses.MeanSquaredError()
optimizer = tf.keras.optimizers.Adam(learning_rate=0.01)
for i in range(10):
tf.summary.experimental.set_step(i)
features = tf.random.normal((16, 32))
labels = tf.random.normal((16, 2))
loss_value = train_step(model, loss, optimizer, features, labels)
print(loss_value)
tf.summary.trace_export("model", profiler_outdir="model")
if __name__ == "__main__":
train()
这没有正确显示模型图，在做
tensorboard --logdir model
在我看到的图表选项卡中
当我通过 model.fit 或 estimator 进行训练时，我得到了图表。例如，这是我model_to_estimator用来转换模型时的图表部分
指南文章没有通过 tensorboard 跟踪指标，我也没有找到任何关于在 alpha 版 ( https://www.tensorflow.org/alpha ) 中自定义添加和跟踪指标的新工作流程的部分。我设计的实现基于 tf.summary 的 API 文档（https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/summary）