I am trying custom training on TensorFlow 2.0 alpha and at the same time I am trying to add some metrics and my training graph to TensorBoard. Consider the following contrived example
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()
This, does not show the model graph properly, on doing
tensorboard --logdir model
In the graphs tab I am seeing 
I am getting the graph when I am training through model.fit or estimator. For example, here is the graphs section when I use model_to_estimator to convert a model
The guide article does not track metrics through tensorboard, and I did not find any sections on the new workflow for custom adding and tracking of metrics in TensorBoard on alpha (https://www.tensorflow.org/alpha). My contrived implementation is based on the API documentation of tf.summary (https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/summary)
