TensorFlow 2.0 Keras：如何为TensorBoard编写图像摘要

Posted 2021-03-31

我正在尝试使用TensorFlow 2.0设置图像识别CNN。为了能够分析我的图像增强，我想看到我在张量板中输入网络的图像。

不幸的是，我无法弄清楚，如何使用TensorFlow 2.0和Keras做到这一点。我也没有找到关于此的文档。

为简单起见，我展示了MNIST示例的代码。我如何在此处添加图像摘要？

import tensorflow as tf
(x_train, y_train), _ = tf.keras.datasets.mnist.load_data()

def scale(image, label):
    return tf.cast(image, tf.float32) / 255.0, label

def augment(image, label):
    return image, label  # do nothing atm

dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
dataset = dataset.map(scale).map(augment).batch(32)

model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(10, activation='softmax')
])

model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
model.fit(dataset, epochs=5, callbacks=[tf.keras.callbacks.TensorBoard(log_dir='D:\tmp\test')])

答案

除了提供你的问题的答案，我将使代码更像TF2.0。如果您有任何疑问/需要澄清，请在下面发表评论。

1. Loading data

我建议使用Tensorflow Datasets库。绝对没有必要在numpy中加载数据并将其转换为tf.data.Dataset，如果可以在一行中完成：

import tensorflow_datasets as tfds

dataset = tfds.load("mnist", as_supervised=True, split=tfds.Split.TRAIN)

上面的行只返回TRAIN分裂（阅读更多关于那些here）。

2. Define Augmentations and Summaries

为了保存图像，必须在每次传递中保持tf.summar.SummaryWriter对象。

我用__call__方法创建了一个方便的包装类，以便于使用tf.data.Dataset的map功能：

import tensorflow as tf

class ExampleAugmentation:
    def __init__(self, logdir: str, max_images: int, name: str):
        self.file_writer = tf.summary.create_file_writer(logdir)
        self.max_images: int = max_images
        self.name: str = name
        self._counter: int = 0

    def __call__(self, image, label):
        augmented_image = tf.image.random_flip_left_right(
            tf.image.random_flip_up_down(image)
        )
        with self.file_writer.as_default():
            tf.summary.image(
                self.name,
                augmented_image,
                step=self._counter,
                max_outputs=self.max_images,
            )

        self._counter += 1
        return augmented_image, label

name将成为保存图像各部分的名称。您可能会问哪一部分 - 由max_outputs定义的部分。

在image中说__call__将具有形状(32, 28, 28, 1)，其中第一个维度是批次，第二个宽度，第三个高度和最后一个通道（仅在MNIST的情况下，但是在tf.image增强中需要这个维度）。此外，让我们说max_outputs被指定为4。在这种情况下，只会保存批次中的4个第一张图像。默认值为3，因此您可以将其设置为BATCH_SIZE以保存每个图像。

在Tensorboard中，每个图像都是一个单独的样本，您可以在最后进行迭代。

需要_counter所以图像不会被覆盖（我想，不太确定，其他人的澄清会很好）。

重要提示：您可能希望在执行更严重的业务时将此类重命名为ImageSaver，并将扩充移动到单独的仿函数/ lambda函数。我想这足以用于演示目的。

3. Setup global variables

请不要混合函数声明，全局变量，数据加载等（如后加载数据和创建函数）。我知道TF1.0鼓励这种类型的节目，但他们正试图摆脱它，你可能想要跟随这种趋势。

下面我已经定义了一些全局变量，这些变量将在下一部分中使用，我猜是非常不言自明的：

BATCH_SIZE = 32
DATASET_SIZE = 60000
EPOCHS = 5

LOG_DIR = "/logs/images"
AUGMENTATION = ExampleAugmentation(LOG_DIR, max_images=4, name="Images")

4. Dataset augmentation

与你的相似，但有一点点扭曲：

dataset = (
    dataset.map(
        lambda image, label: (
            tf.image.convert_image_dtype(image, dtype=tf.float32),
            label,
        )
    )
    .batch(BATCH_SIZE)
    .map(AUGMENTATION)
    .repeat(EPOCHS)
)

• 由于加载的数据集是生成器，因此需要repeat
• tf.image.convert_image_dtype - 比明确的tf.cast与255的分割更好和更易读的选项（并确保正确的图像格式）
• 在扩充之前完成批处理只是为了演示

5. Define model, compile, train

就像你在你的例子中所做的那样，但我提供了额外的steps_per_epoch，所以fit知道有多少批次构成了一个时代：

model = tf.keras.models.Sequential(
    [
        tf.keras.layers.Flatten(input_shape=(28, 28, 1)),
        tf.keras.layers.Dense(128, activation="relu"),
        tf.keras.layers.Dropout(0.2),
        tf.keras.layers.Dense(10, activation="softmax"),
    ]
)

model.compile(
    optimizer="adam", loss="sparse_categorical_crossentropy", metrics=["accuracy"]
)
model.fit(
    dataset,
    epochs=EPOCHS,
    steps_per_epoch=DATASET_SIZE // BATCH_SIZE,
    callbacks=[tf.keras.callbacks.TensorBoard(log_dir=LOG_DIR)],
)

除了我认为的解释之外，没什么可说的。

6. Run Tensorboard

因为TF2.0可以使用%tensorboard --logdir /logs/images在colab内部进行，所以只想为可能访问此问题的其他人添加此内容。无论你喜欢它，无论如何你知道如何做到这一点。

图像应该在IMAGES内部，并且name命名的每个样本都提供给AUGMENTATION对象。

7. Whole code (to make everyone's life easier)

import tensorflow as tf

import tensorflow_datasets as tfds


class ExampleAugmentation:
    def __init__(self, logdir: str, max_images: int, name: str):
        self.file_writer = tf.summary.create_file_writer(logdir)
        self.max_images: int = max_images
        self.name: str = name
        self._counter: int = 0

    def __call__(self, image, label):
        augmented_image = tf.image.random_flip_left_right(
            tf.image.random_flip_up_down(image)
        )
        with self.file_writer.as_default():
            tf.summary.image(
                self.name,
                augmented_image,
                step=self._counter,
                max_outputs=self.max_images,
            )

        self._counter += 1
        return augmented_image, label


# Global settings

BATCH_SIZE = 32
DATASET_SIZE = 60000
EPOCHS = 5

LOG_DIR = "/logs/images"
AUGMENTATION = ExampleAugmentation(LOG_DIR, max_images=4, name="Images")

# Dataset

dataset = tfds.load("mnist", as_supervised=True, split=tfds.Split.TRAIN)

dataset = (
    dataset.map(
        lambda image, label: (
            tf.image.convert_image_dtype(image, dtype=tf.float32),
            label,
        )
    )
    .batch(BATCH_SIZE)
    .map(AUGMENTATION)
    .repeat(EPOCHS)
)

# Model and training

model = tf.keras.models.Sequential(
    [
        tf.keras.layers.Flatten(input_shape=(28, 28, 1)),
        tf.keras.layers.Dense(128, activation="relu"),
        tf.keras.layers.Dropout(0.2),
        tf.keras.layers.Dense(10, activation="softmax"),
    ]
)

model.compile(
    optimizer="adam", loss="sparse_categorical_crossentropy", metrics=["accuracy"]
)
model.fit(
    dataset,
    epochs=EPOCHS,
    steps_per_epoch=DATASET_SIZE // BATCH_SIZE,
    callbacks=[tf.keras.callbacks.TensorBoard(log_dir=LOG_DIR)],
)

另一答案

你可以做这样的事情来将输入图像添加到张量板

def scale(image, label):
    return tf.cast(image, tf.float32) / 255.0, label


def augment(image, label):
    return image, label  # do nothing atm


file_writer = tf.summary.create_file_writer(logdir + "/images")


def plot_to_image(figure):
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    plt.close(figure)
    buf.seek(0)
    image = tf.image.decode_png(buf.getvalue(), channels=4)
    image = tf.expand_dims(image, 0)
    return image


def image_grid():
    """Return a 5x5 grid of the MNIST images as a matplotlib figure."""
    # Create a figure to contain the plot.
    figure = plt.figure(figsize=(10, 10))
    for i in range(25):
        # Start next subplot.
        plt.subplot(5, 5, i + 1, title=str(y_train[i]))
        plt.xticks([])
        plt.yticks([])
        plt.grid(False)
        image, _ = scale(x_train[i], y_train[i])
        plt.imshow(x_train[i], cmap=plt.cm.binary)

    return figure


# Prepare the plot
figure = image_grid()
# Convert to image and log
with file_writer.as_default():
    tf.summary.image("Training data", plot_to_image(figure), step=0)

dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
dataset = dataset.map(scale).map(augment).batch(32)

model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(10, activation='softmax')
])

model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

model.fit(dataset, epochs=5, callbacks=[tf.keras.callbacks.TensorBoard(log_dir=logdir)])