除了回答您的问题
我会让代码更像TF2.0-like。如果您有任何问题/需要澄清，请在下方发表评论。
1。加载数据
我建议使用Tensorflow Datasets 库。完全不需要在numpy 中加载数据并将其转换为tf.data.Dataset，如果可以在一行中完成：
import tensorflow_datasets as tfds
dataset = tfds.load("mnist", as_supervised=True, split=tfds.Split.TRAIN)
上面的行只会返回TRAIN split（阅读更多关于here的信息）。
2。定义扩充和总结
为了保存图像，必须在每次传递中保留tf.summary.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 定义的部分。
在__call__ 中说image 将具有(32, 28, 28, 1) 的形状，其中第一个维度是批次、第二个宽度、第三个高度和最后一个通道（在MNIST 的情况下只有一个，但在tf.image 增强中需要此维度）。此外，假设max_outputs 被指定为4。在这种情况下，批次中的前 4 张图像将被保存。默认值为3，因此您可以将其设置为BATCH_SIZE 以保存每张图像。
在Tensorboard 中，每个图像都是一个单独的样本，您可以在最后对其进行迭代。
_counter 是必需的，因此图像不会被覆盖（我认为，不太确定，请其他人澄清一下）。
重要提示：在进行更严肃的业务时，您可能希望将此类重命名为 ImageSaver，并将扩充移动到单独的仿函数/lambda 函数。我猜它足以用于演示目的。
3。设置全局变量
请不要混用函数声明、全局变量、数据加载等（如加载数据和事后创建函数）。我知道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 = (
dataset.map(
lambda image, label: (
tf.image.convert_image_dtype(image, dtype=tf.float32),
label,
)
)
.batch(BATCH_SIZE)
.map(AUGMENTATION)
.repeat(EPOCHS)
)
5。定义模型、编译、训练
几乎与您在示例中所做的一样，但我提供了额外的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。运行 Tensorboard
由于 TF2.0 可以在 colab 中使用 %tensorboard --logdir /logs/images 完成，只是想为可能访问此问题的其他人添加此内容。随心所欲地去做，反正你肯定知道怎么做。
图像应位于IMAGES 内部，并且每个样本均以name 命名，并提供给AUGMENTATION 对象。
7.完整的代码（让每个人的生活更轻松）
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
if __name__ == "__main__":
# 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)],
)