[Python, Tensorflow] Tensorboard 내 이미지 추가 방법 (Classficiation)

안녕하세요, 오늘은 Tensorboard 내 학습 중 사용하는 Images를 추가하는 방법에 대해서 알아보겠습니다.

 

전체 코드는 다음과 같습니다.

 

import tensorflow as tf
from tensorflow import keras
import numpy as np
import datetime

class CustomImageLogger(tf.keras.callbacks.Callback):
    def __init__(self, log_dir, validation_data, num_images=3):
        super().__init__()
        self.log_dir = log_dir
        self.num_images = num_images
        self.validation_data = validation_data

    def on_epoch_end(self, epoch, logs=None):
        file_writer = tf.summary.create_file_writer(self.log_dir + '/img')

        with file_writer.as_default():

            for i in range(self.num_images):
                img = self.validation_data[0][i] 
                img = tf.expand_dims(img, 0) 
                img = tf.expand_dims(img, -1)
                
                tf.summary.image("Validation Image at epoch {} - Image number {}".format(epoch, i), img, step=epoch)

(train_images, train_labels), (valid_images, valid_labels) = tf.keras.datasets.mnist.load_data()
train_images = train_images / 255.0
valid_images = valid_images / 255.0

model = ...

log_dir = "logs/model"  + "/" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=1)
image_logger_callback = CustomImageLogger(log_dir=log_dir, validation_data=(valid_images, valid_labels))


model.fit(train_images, train_labels, epochs=100, validation_data=(valid_images, valid_labels),
          callbacks=[tensorboard_callback, image_logger_callback])

 

위의 전체 코드에서 주목해야할 부분은 다음과 같습니다.

 

class CustomImageLogger(tf.keras.callbacks.Callback):
    def __init__(self, log_dir, validation_data, num_images=3):
        super().__init__()
        self.log_dir = log_dir
        self.num_images = num_images
        self.validation_data = validation_data

    def on_epoch_end(self, epoch, logs=None):
        file_writer = tf.summary.create_file_writer(self.log_dir + '/img')

        with file_writer.as_default():

            for i in range(self.num_images):
                img = self.validation_data[0][i] 
                img = tf.expand_dims(img, 0) 
                img = tf.expand_dims(img, -1)
                
                tf.summary.image("Validation Image at epoch {} - Image number {}".format(epoch, i), img, step=epoch)

 

위의 코드는 각 Epoch 종료될 때 마다  Validation Images 내 3장의 이미지를 각 Epoch 당 Log에 기록하게 됩니다.

 

Tensorboard에서 실행한  결과는 다음과 같습니다.

 

Tensorboard 내 이미지 추가된 결과

 

Tensorboard 내 "IMAGES" 메뉴가 새로 생긴 것을 확인할 수 있으며, 각 Epoch 당 3개의 이미지가 추가된 모습을 확인할 수 있습니다.

 

 

지금까지 ensorboard 내 학습 중 사용하는 Images를 추가하는 방법에 대해서 알아보았습니다.

감사드립니다.