"warning:tensorflow:`write_grads` will be ignored in tensorflow 2.0 for the `tensorboard` callback" Code Answer

write_grads was not implemented in tf2.x. this is one of the highly expected feature request that is still open. please check this github issue as feature request. so, we only need to import tf1.x modules and use write_grads as shown in the following code.

# load the tensorboard notebook extension
%load_ext tensorboard

import tensorflow as tf
import datetime

# clear any logs from previous runs
!rm -rf ./logs/ 

# disable v2 behavior
tf.compat.v1.disable_v2_behavior()

mnist = tf.keras.datasets.mnist

(x_train, y_train),(x_test, y_test) = mnist.load_data()

x_train, x_test = x_train / 255.0, x_test / 255.0


def create_model():

  return tf.keras.models.sequential([

    tf.keras.layers.flatten(input_shape=(28, 28)),

    tf.keras.layers.dense(512, activation='relu'),

    tf.keras.layers.dropout(0.2),

    tf.keras.layers.dense(10, activation='softmax')

  ])

 

model = create_model()

model.compile(optimizer='adam',

              loss='sparse_categorical_crossentropy',

              metrics=['accuracy'])


log_dir = "logs/fit/" + datetime.datetime.now().strftime("%y%m%d-%h%m%s")

tensorboard_callback = tf.compat.v1.keras.callbacks.tensorboard(log_dir=log_dir, histogram_freq=1, write_grads =true)

model.fit(x=x_train, y=y_train, epochs=1, validation_data=(x_test, y_test), callbacks=[tensorboard_callback]) 

%tensorboard --logdir logs/fit

output:

downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz
11493376/11490434 [==============================] - 0s 0us/step

train on 60000 samples, validate on 10000 samples
warning:tensorflow:from /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training_v1.py:2048: model.state_updates (from tensorflow.python.keras.engine.training) is deprecated and will be removed in a future version.
instructions for updating:
this property should not be used in tensorflow 2.0, as updates are applied automatically.
   32/60000 [..............................] - eta: 0s - loss: 2.3311 - acc: 0.0312warning:tensorflow:callbacks method `on_train_batch_end` is slow compared to the batch time (batch time: 0.0055s vs `on_train_batch_end` time: 0.0235s). check your callbacks.
60000/60000 [==============================] - 17s 288us/sample - loss: 0.2187 - acc: 0.9349 - val_loss: 0.1012 - val_acc: 0.9690
<tensorflow.python.keras.callbacks.history at 0x7f7ebd1d3d30>