在Keras下使用Tensorboard写入图片

在Keras中调用Tensorboard可以直接创建一个Tensorboard对象，在model.fit的callbacks里面就好了，但是有的时候需要添加图片到Tensorboard里面去，直接使用Tensorboard就不行了，需要重写一下Tensorboard。比如在下图中，需要将验证数据集的第一个数据绘制出来并显示结果，那么

import keras
import tensorflow as tf
import matplotlib.pyplot as plt
import io

class CustomTensorBoard(keras.callbacks.TensorBoard):

    def _make_image(self):
        buf = io.BytesIO()

        y = self.model.predict(self.validation_data[0][0:1])
        plt.plot(self.validation_data[0][0])
        plt.title(f'Predict Result:{y}')
        plt.savefig(buf, format='png')
        plt.close()

        buf.seek(0)
        image = buf.getvalue()
        image = tf.Summary.Image(encoded_image_string=image)
        return image

    def on_epoch_end(self, epoch, logs=None):
        super(CustomTensorBoard, self).on_epoch_end(epoch, logs)
        image = self._make_image()
        summary = tf.Summary(value=[tf.Summary.Value(tag='plot', image=image)])
        self.writer.add_summary(summary, global_step=epoch)
		
		# 如果需要添加其他数字
		# summary = tf.Summary(value=[tf.Summary.Value(tag='my_number', simple_value=my_number)])
        # self.writer.add_summary(summary, global_step=epoch)
	

ctb = CustomTensorBoard('./logs')

model=Model(......)
model.fit(...,callbacks=[ctb])

通过_make_image生成一张图片，也可以通过其他的方式生成图片。然后重写on_epoch_end这个方法，super(CustomTensorBoard, self).on_epoch_end(epoch, logs)先调用父类方法，然后在写自己的代码，可以在不改变原来功能的情况下添加新的功能。

另外找到一种方式，但绘制的图片不是按tag分组的，但是也可适用的。

import keras
import tensorflow as tf
import matplotlib.pyplot as plt
import io

class CustomTensorBoard(keras.callbacks.TensorBoard):

    def _make_image(self):
        buf = io.BytesIO()
        y = self.model.predict(self.validation_data[0][0:1])
        plt.plot(self.validation_data[0][0])
        plt.title(f'Predict Result:{np.argmax(y)}')
        plt.savefig(buf, format='png')
        plt.close()

        buf.seek(0)
        image = tf.image.decode_png(buf.getvalue(), channels=4)
        image = tf.expand_dims(image, 0)
        return image

    def on_epoch_end(self, epoch, logs=None):
        super(CustomTensorBoard, self).on_epoch_end(epoch, logs)
        image = self._make_image()
        with tf.Session() as sess:
            summary_op = tf.summary.image('plot', image)
            summary = sess.run(summary_op)
            self.writer.add_summary(summary, global_step=epoch)

ctb = CustomTensorBoard('./logs')

model=Model(......)
model.fit(...,callbacks=[ctb])