tensorflow2.x 如何快捷的自定义使用tensorboard

前言

tensorflow2.x常常于keras配合使用，官方提供了两种tensorboard使用方案，一种是使用keras.callbacks,另一种则是tf.summary自定义。第一种方案可以简单地得到计算图，模型权重，处理时间分配等信息，第二种可以想画什么画什么但必须自定义训练过程无法直接在fit中使用。有没有一种简单方式可以自定义tensorboard又不用自定义训练过程呢。

了解callbacks

keras.callbacks.Callback是提供给fit函数的回调函数（类）基类，所有的callbacks都继承了这个类。
这个类有如下常用属性和方法：

set_model(self, model)
on_batch_begin(self, batch, logs=None)
on_batch_end(self, batch, logs=None)
on_epoch_begin(self, epoch, logs=None)
on_epoch_end(self, epoch, logs=None)
on_train_batch_begin(self, batch, logs=None)
on_train_batch_end(self, batch, logs=None)
on_test_batch_begin(self, batch, logs=None)
on_test_batch_end(self, batch, logs=None)
on_predict_batch_begin(self, batch, logs=None)
on_predict_batch_end(self, batch, logs=None)
on_train_begin(self, logs=None)
on_train_end(self, logs=None)
on_test_begin(self, logs=None)
on_test_end(self, logs=None)
on_predict_begin(self, logs=None)
on_predict_end(self, logs=None)
调用的时间和参数显而易见（忽略logs参数）

自定义Tensorboard的callback

示例代码

import tensorflow as tf
from tensorflow import keras

mnist = 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
x_train=x_train[...,None]#60000,28,28 变成 60000,28,28，1
x_test=x_test[...,None]

class myTensorboard(keras.callbacks.TensorBoard):
    def __init__(self,validation_x,show_num=10, **kwargs):
        super().__init__(**kwargs)
        self.show_num=show_num # 要显示的图片数量
        self.validation_x=validation_x
    def on_epoch_end(self, epoch, logs=None):
        super(myTensorboard, self).on_epoch_end(epoch,logs)
        writer = self._get_writer(self._validation_run_name)
        x=tf.random.shuffle(self.validation_x)[:self.show_num]
        conv1=self.model.conv1(x)
        y=self.model(x)
        with writer.as_default():
            for i,(x,y) in enumerate(zip(x,y)):
                tf.summary.image(f'img/{i}',data=x[None,...],step=epoch)
                tf.summary.image(f'label/{i}',data=tf.reshape(y,(1,1,10,1)),step=epoch)
                for j in range(32):#输出每个卷积核卷积结果
                    tf.summary.image(f'conv/{i}_{j}',data=conv1[...,j:j+1],step=epoch)

        tf.summary.flush(writer)


class simpleCnn(keras.Model):
    def build(self, input_shape):
        super(simpleCnn, self).build(input_shape)
        self.conv1=keras.layers.Conv2D(32,3,padding='same',
                                       activation=keras.activations.relu,
                                       name='conv1/conv',
                                       kernel_regularizer=keras.regularizers.l1(),bias_regularizer=keras.regularizers.l1())
        self.pool1=keras.layers.MaxPool2D(padding='same',name='conv1/pool')
        self.conv2=keras.layers.Conv2D(32,3,padding='same',activation=keras.activations.relu,name='conv2/conv')
        self.pool2=keras.layers.MaxPool2D(padding='same',name='conv2/pool')
        self.fc=keras.layers.Flatten()
        self.ol=keras.layers.Dense(10,activation=keras.activations.softmax,name='output')
    def call(self, inputs, training=None, mask=None):
        x=inputs
        x=self.conv1(x)
        x=self.pool1(x)
        x=self.conv2(x)
        x=self.pool2(x)
        x=self.fc(x)
        x=self.ol(x)
        return x

myboard=myTensorboard(x_test)
model=simpleCnn()
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])
model.fit(x_train, y_train, epochs=5,validation_data=(x_test,y_test),batch_size=128,callbacks=[myboard])

效果

这个例子中自定义了图片的显示，实际上你也可以用类似的方式自定义直方图，散点图等。
还有几个类似代码中tf.summary的函数库，如tensorboard.summary,tensorflow.python.ops.summary_ops_v2，使用方式上大体一样，目前还不明白有什么区别，为什么要分这么多类似的库。

其他的一些小技巧

我常常在使用keras搭建一些输入输出很复杂的模型时想用tensorboard检查一下，这时只需要用keras.callbacks.Tensorboard().setmodel(model)就可以马上看到已经搭好的计算图，是不是比直接用tensorflow提供的summary轻松很多呢。

总结

找不到文档时不妨看看源码。