使用功能子类化对Keras模型进行广泛分发的打包

import keras
import tensorflow as tf  # 仅用于 tf.data

print("Keras版本", keras.version())
print("Keras运行在", keras.config.backend())

# tf.data是一个用于组合数据流的优秀API。
# 无论您使用TensorFlow、PyTorch还是Jax后端，它都能工作，
# 只要您在数据流中使用它，而不是在模型内部。

BATCH_SIZE = 256

(x_train, train_labels), (x_test, test_labels) = keras.datasets.mnist.load_data()

train_data = tf.data.Dataset.from_tensor_slices((x_train, train_labels))
train_data = train_data.map(
    lambda x, y: (tf.expand_dims(x, axis=-1), y)
)  # 1通道单色
train_data = train_data.batch(BATCH_SIZE)
train_data = train_data.cache()
train_data = train_data.shuffle(5000, reshuffle_each_iteration=True)
train_data = train_data.repeat()

test_data = tf.data.Dataset.from_tensor_slices((x_test, test_labels))
test_data = test_data.map(
    lambda x, y: (tf.expand_dims(x, axis=-1), y)
)  # 1通道单色
test_data = test_data.batch(10000)
test_data = test_data.cache()

STEPS_PER_EPOCH = len(train_labels) // BATCH_SIZE
EPOCHS = 5

class MnistModel(keras.Model):
    def __init__(self, **kwargs):
        # Keras函数模型的定义。这也可以使用Sequential，
        # Sequential只是简单函数模型的语法糖。

        # 1通道单色输入
        inputs = keras.layers.Input(shape=(None, None, 1), dtype="uint8")
        # 像素格式从uint8转换为float32
        y = keras.layers.Rescaling(1 / 255.0)(inputs)

        # 3个卷积层
        y = keras.layers.Conv2D(
            filters=16, kernel_size=3, padding="same", activation="relu"
        )(y)
        y = keras.layers.Conv2D(
            filters=32, kernel_size=6, padding="same", activation="relu", strides=2
        )(y)
        y = keras.layers.Conv2D(
            filters=48, kernel_size=6, padding="same", activation="relu", strides=2
        )(y)

        # 2个全连接层
        y = keras.layers.GlobalAveragePooling2D()(y)
        y = keras.layers.Dense(48, activation="relu")(y)
        y = keras.layers.Dropout(0.4)(y)
        outputs = keras.layers.Dense(
            10, activation="softmax", name="classification_head"  # 10个类别
        )(y)

        # 通过调用keras.Model(inputs, outputs)创建Keras函数模型
        super().__init__(inputs=inputs, outputs=outputs, **kwargs)

model = MnistModel()

model.compile(
    optimizer="adam",
    loss="sparse_categorical_crossentropy",
    metrics=["sparse_categorical_accuracy"],
)

history = model.fit(
    train_data,
    steps_per_epoch=STEPS_PER_EPOCH,
    epochs=EPOCHS,
    validation_data=test_data,
)

model = MnistModel()

# 模型摘要有效
model.summary()


# 递归遍历层图也有效
def walk_layers(layer):
    if hasattr(layer, "layers"):
        for layer in layer.layers:
            walk_layers(layer)
    else:
        print(layer.name)


print("\n遍历模型层：\n")
walk_layers(model)

model = MnistModel()

input = model.input
# 在分类头之前截取
y = model.get_layer("classification_head").input

# 添加新的分类头
output = keras.layers.Dense(
    1,  # 二分类的单一类别
    activation="sigmoid",
    name="binary_classification_head",
)(y)

# 创建一个新的函数模型
binary_model = keras.Model(input, output)

binary_model.summary()

# 新数据集，标签为0 / 1（1 = 数字 '0'，0 = 所有其他数字）
bin_train_data = train_data.map(
    lambda x, y: (x, tf.cast(tf.math.equal(y, tf.zeros_like(y)), dtype=tf.uint8))
)
bin_test_data = test_data.map(
    lambda x, y: (x, tf.cast(tf.math.equal(y, tf.zeros_like(y)), dtype=tf.uint8))
)

# 对于二分类的合适损失和评估指标
binary_model.compile(
    optimizer="adam", loss="binary_crossentropy", metrics=["binary_accuracy"]
)

history = binary_model.fit(
    bin_train_data,
    steps_per_epoch=STEPS_PER_EPOCH,
    epochs=EPOCHS,
    validation_data=bin_test_data,
)

class MnistDictModel(keras.Model):
    def __init__(self, **kwargs):
        #
        # 输入为字典
        #
        inputs = {
            "image": keras.layers.Input(
                shape=(None, None, 1),  # 1通道单色
                dtype="uint8",
                name="image",
            )
        }

        # 像素格式从uint8转换为float32
        y = keras.layers.Rescaling(1 / 255.0)(inputs["image"])

        # 3个卷积层
        y = keras.layers.Conv2D(
            filters=16, kernel_size=3, padding="same", activation="relu"
        )(y)
        y = keras.layers.Conv2D(
            filters=32, kernel_size=6, padding="same", activation="relu", strides=2
        )(y)
        y = keras.layers.Conv2D(
            filters=48, kernel_size=6, padding="same", activation="relu", strides=2
        )(y)

        # 2个全连接层
        y = keras.layers.GlobalAveragePooling2D()(y)
        y = keras.layers.Dense(48, activation="relu")(y)
        y = keras.layers.Dropout(0.4)(y)
        outputs = keras.layers.Dense(
            10, activation="softmax", name="classification_head"  # 10个类别
        )(y)

        # 通过调用 keras.Model(inputs, outputs) 创建 Keras 功能模型
        super().__init__(inputs=inputs, outputs=outputs, **kwargs)

model = MnistDictModel()

# 将数据集重新格式化为字典
dict_train_data = train_data.map(lambda x, y: ({"image": x}, y))
dict_test_data = test_data.map(lambda x, y: ({"image": x}, y))

model.compile(
    optimizer="adam",
    loss="sparse_categorical_crossentropy",
    metrics=["sparse_categorical_accuracy"],
)

history = model.fit(
    dict_train_data,
    steps_per_epoch=STEPS_PER_EPOCH,
    epochs=EPOCHS,
    validation_data=dict_test_data,
)