没有注意力机制的视觉变换器

!pip install -qq -U tensorflow-addons

import numpy as np
import matplotlib.pyplot as plt

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import tensorflow_addons as tfa

import pathlib
import glob

# 为了重现性设置随机种子
SEED = 42
keras.utils.set_random_seed(SEED)

class Config(object):
    # 数据
    batch_size = 256
    buffer_size = batch_size * 2
    input_shape = (32, 32, 3)
    num_classes = 10

    # 增强
    image_size = 48

    # 架构
    patch_size = 4
    projected_dim = 96
    num_shift_blocks_per_stages = [2, 4, 8, 2]
    epsilon = 1e-5
    stochastic_depth_rate = 0.2
    mlp_dropout_rate = 0.2
    num_div = 12
    shift_pixel = 1
    mlp_expand_ratio = 2

    # 优化器
    lr_start = 1e-5
    lr_max = 1e-3
    weight_decay = 1e-4

    # 训练
    epochs = 100

    # 推理
    label_map = {
        0: "飞机",
        1: "汽车",
        2: "鸟",
        3: "猫",
        4: "鹿",
        5: "狗",
        6: "青蛙",
        7: "马",
        8: "船",
        9: "卡车",
    }
    tf_ds_batch_size = 20


config = Config()

(x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()
(x_train, y_train), (x_val, y_val) = (
    (x_train[:40000], y_train[:40000]),
    (x_train[40000:], y_train[40000:]),
)
print(f"训练样本: {len(x_train)}")
print(f"验证样本: {len(x_val)}")
print(f"测试样本: {len(x_test)}")

AUTO = tf.data.AUTOTUNE
train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_ds = train_ds.shuffle(config.buffer_size).batch(config.batch_size).prefetch(AUTO)

val_ds = tf.data.Dataset.from_tensor_slices((x_val, y_val))
val_ds = val_ds.batch(config.batch_size).prefetch(AUTO)

test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test))
test_ds = test_ds.batch(config.batch_size).prefetch(AUTO)

def get_augmentation_model():
    """构建数据增强模型。"""
    data_augmentation = keras.Sequential(
        [
            layers.Resizing(config.input_shape[0] + 20, config.input_shape[0] + 20),
            layers.RandomCrop(config.image_size, config.image_size),
            layers.RandomFlip("horizontal"),
            layers.Rescaling(1 / 255.0),
        ]
    )
    return data_augmentation

class MLP(layers.Layer):
    """获取每个移位块的MLP层。

    参数：
        mlp_expand_ratio (int)：第一个特征图扩展的比率。
        mlp_dropout_rate (float)：丢弃的比率。
    """

    def __init__(self, mlp_expand_ratio, mlp_dropout_rate, **kwargs):
        super().__init__(**kwargs)
        self.mlp_expand_ratio = mlp_expand_ratio
        self.mlp_dropout_rate = mlp_dropout_rate

    def build(self, input_shape):
        input_channels = input_shape[-1]
        initial_filters = int(self.mlp_expand_ratio * input_channels)

        self.mlp = keras.Sequential(
            [
                layers.Dense(
                    units=initial_filters,
                    activation=tf.nn.gelu,
                ),
                layers.Dropout(rate=self.mlp_dropout_rate),
                layers.Dense(units=input_channels),
                layers.Dropout(rate=self.mlp_dropout_rate),
            ]
        )

    def call(self, x):
        x = self.mlp(x)
        return x

class DropPath(layers.Layer):
    """Drop Path，也称为随机深度层。

    参考：
        - https://keras.io/examples/vision/cct/#stochastic-depth-for-regularization
        - github.com:rwightman/pytorch-image-models
    """

    def __init__(self, drop_path_prob, **kwargs):
        super().__init__(**kwargs)
        self.drop_path_prob = drop_path_prob

    def call(self, x, training=False):
        if training:
            keep_prob = 1 - self.drop_path_prob
            shape = (tf.shape(x)[0],) + (1,) * (len(tf.shape(x)) - 1)
            random_tensor = keep_prob + tf.random.uniform(shape, 0, 1)
            random_tensor = tf.floor(random_tensor)
            return (x / keep_prob) * random_tensor
        return x

out[:, g * 0:g * 1, :, :-1] = x[:, g * 0:g * 1, :, 1:]  # 向左移动
out[:, g * 1:g * 2, :, 1:] = x[:, g * 1:g * 2, :, :-1]  # 向右移动
out[:, g * 2:g * 3, :-1, :] = x[:, g * 2:g * 3, 1:, :]  # 向上移动
out[:, g * 3:g * 4, 1:, :] = x[:, g * 3:g * 4, :-1, :]  # 向下移动

out[:, g * 4:, :, :] = x[:, g * 4:, :, :]  # 不移动

class ShiftViTBlock(layers.Layer):
    """一个单位 ShiftViT 块

    参数：
        shift_pixel (int): 移动的像素数。默认为1。
        mlp_expand_ratio (int): MLP特征扩展的比例。
            默认为2。
        mlp_dropout_rate (float): MLP中使用的丢弃率。
        num_div (int): 特征图通道的分割数。
            总共将移动 4/num_div 的通道。默认为12。
        epsilon (float): Epsilon 常量。
        drop_path_prob (float): 丢弃路径的丢弃概率。
    """

    def __init__(
        self,
        epsilon,
        drop_path_prob,
        mlp_dropout_rate,
        num_div=12,
        shift_pixel=1,
        mlp_expand_ratio=2,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.shift_pixel = shift_pixel
        self.mlp_expand_ratio = mlp_expand_ratio
        self.mlp_dropout_rate = mlp_dropout_rate
        self.num_div = num_div
        self.epsilon = epsilon
        self.drop_path_prob = drop_path_prob

    def build(self, input_shape):
        self.H = input_shape[1]
        self.W = input_shape[2]
        self.C = input_shape[3]
        self.layer_norm = layers.LayerNormalization(epsilon=self.epsilon)
        self.drop_path = (
            DropPath(drop_path_prob=self.drop_path_prob)
            if self.drop_path_prob > 0.0
            else layers.Activation("linear")
        )
        self.mlp = MLP(
            mlp_expand_ratio=self.mlp_expand_ratio,
            mlp_dropout_rate=self.mlp_dropout_rate,
        )

    def get_shift_pad(self, x, mode):
        """根据所选的模式移动通道。"""
        if mode == "left":
            offset_height = 0
            offset_width = 0
            target_height = 0
            target_width = self.shift_pixel
        elif mode == "right":
            offset_height = 0
            offset_width = self.shift_pixel
            target_height = 0
            target_width = self.shift_pixel
        elif mode == "up":
            offset_height = 0
            offset_width = 0
            target_height = self.shift_pixel
            target_width = 0
        else:
            offset_height = self.shift_pixel
            offset_width = 0
            target_height = self.shift_pixel
            target_width = 0
        crop = tf.image.crop_to_bounding_box(
            x,
            offset_height=offset_height,
            offset_width=offset_width,
            target_height=self.H - target_height,
            target_width=self.W - target_width,
        )
        shift_pad = tf.image.pad_to_bounding_box(
            crop,
            offset_height=offset_height,
            offset_width=offset_width,
            target_height=self.H,
            target_width=self.W,
        )
        return shift_pad

    def call(self, x, training=False):
        # 分割特征图
        x_splits = tf.split(x, num_or_size_splits=self.C // self.num_div, axis=-1)

        # 移动特征图
        x_splits[0] = self.get_shift_pad(x_splits[0], mode="left")
        x_splits[1] = self.get_shift_pad(x_splits[1], mode="right")
        x_splits[2] = self.get_shift_pad(x_splits[2], mode="up")
        x_splits[3] = self.get_shift_pad(x_splits[3], mode="down")

        # 连接移动和未移动的特征图
        x = tf.concat(x_splits, axis=-1)

        # 添加残差连接
        shortcut = x
        x = shortcut + self.drop_path(self.mlp(self.layer_norm(x)), training=training)
        return x

class PatchMerging(layers.Layer):
    """补丁合并层。

    Args:
        epsilon (float): epsilon常数。
    """

    def __init__(self, epsilon, **kwargs):
        super().__init__(**kwargs)
        self.epsilon = epsilon

    def build(self, input_shape):
        filters = 2 * input_shape[-1]
        self.reduction = layers.Conv2D(
            filters=filters, kernel_size=2, strides=2, padding="same", use_bias=False
        )
        self.layer_norm = layers.LayerNormalization(epsilon=self.epsilon)

    def call(self, x):
        # 在特征图上应用补丁合并算法
        x = self.layer_norm(x)
        x = self.reduction(x)
        return x

# 注意：这一层在模型的不同阶段会有不同的堆叠深度
class StackedShiftBlocks(layers.Layer):
    """包含堆叠的 ShiftViTBlocks 的层。

    参数：
        epsilon (float): epsilon 常数。
        mlp_dropout_rate (float): MLP 块中使用的 dropout 率。
        num_shift_blocks (int): 此阶段的移位 vit 块数量。
        stochastic_depth_rate (float): 选择的最大 dropout 路径率。
        is_merge (boolean): 一个标志，用于确定是否在移位 vit 块之后使用 Patch Merge 层。
        num_div (int): 特征图的通道数划分。默认为 12。
        shift_pixel (int): 移位的像素数。默认为 1。
        mlp_expand_ratio (int): MLP 的初始密集层扩展的比例。默认为 2。
    """

    def __init__(
        self,
        epsilon,
        mlp_dropout_rate,
        num_shift_blocks,
        stochastic_depth_rate,
        is_merge,
        num_div=12,
        shift_pixel=1,
        mlp_expand_ratio=2,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.epsilon = epsilon
        self.mlp_dropout_rate = mlp_dropout_rate
        self.num_shift_blocks = num_shift_blocks
        self.stochastic_depth_rate = stochastic_depth_rate
        self.is_merge = is_merge
        self.num_div = num_div
        self.shift_pixel = shift_pixel
        self.mlp_expand_ratio = mlp_expand_ratio

    def build(self, input_shapes):
        # 计算随机深度概率。
        # 参考: https://keras.io/examples/vision/cct/#the-final-cct-model
        dpr = [
            x
            for x in np.linspace(
                start=0, stop=self.stochastic_depth_rate, num=self.num_shift_blocks
            )
        ]

        # 将移位块构建为 ShiftViT 块的列表
        self.shift_blocks = list()
        for num in range(self.num_shift_blocks):
            self.shift_blocks.append(
                ShiftViTBlock(
                    num_div=self.num_div,
                    epsilon=self.epsilon,
                    drop_path_prob=dpr[num],
                    mlp_dropout_rate=self.mlp_dropout_rate,
                    shift_pixel=self.shift_pixel,
                    mlp_expand_ratio=self.mlp_expand_ratio,
                )
            )
        if self.is_merge:
            self.patch_merge = PatchMerging(epsilon=self.epsilon)

    def call(self, x, training=False):
        for shift_block in self.shift_blocks:
            x = shift_block(x, training=training)
        if self.is_merge:
            x = self.patch_merge(x)
        return x

    # 由于这是一个自定义层，我们需要重写 get_config()
    # 以便模型可以在训练后方便地保存和加载
    def get_config(self):
        config = super().get_config()
        config.update(
            {
                "epsilon": self.epsilon,
                "mlp_dropout_rate": self.mlp_dropout_rate,
                "num_shift_blocks": self.num_shift_blocks,
                "stochastic_depth_rate": self.stochastic_depth_rate,
                "is_merge": self.is_merge,
                "num_div": self.num_div,
                "shift_pixel": self.shift_pixel,
                "mlp_expand_ratio": self.mlp_expand_ratio,
            }
        )
        return config

class ShiftViTModel(keras.Model):
    """ShiftViT模型。

    参数：
        data_augmentation (keras.Model): 数据增强模型。
        projected_dim (int): 图像块投影的维度。
        patch_size (int): 图像的块大小。
        num_shift_blocks_per_stages (list[int]): 每个阶段的所有移位块的数量列表。
        epsilon (float): epsilon常数。
        mlp_dropout_rate (float): 在MLP块中使用的丢弃率。
        stochastic_depth_rate (float): 最大丢弃率概率。
        num_div (int): 特征图通道的划分数量。默认为12。
        shift_pixel (int): 移位的像素数量。默认为1。
        mlp_expand_ratio (int): 初始MLP密集层扩展到的比例。默认为2。
    """


    def __init__(
        self,
        data_augmentation,
        projected_dim,
        patch_size,
        num_shift_blocks_per_stages,
        epsilon,
        mlp_dropout_rate,
        stochastic_depth_rate,
        num_div=12,
        shift_pixel=1,
        mlp_expand_ratio=2,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.data_augmentation = data_augmentation
        self.patch_projection = layers.Conv2D(
            filters=projected_dim,
            kernel_size=patch_size,
            strides=patch_size,
            padding="same",
        )
        self.stages = list()
        for index, num_shift_blocks in enumerate(num_shift_blocks_per_stages):
            if index == len(num_shift_blocks_per_stages) - 1:
                # This is the last stage, do not use the patch merge here.
                is_merge = False
            else:
                is_merge = True
            # Build the stages.
            self.stages.append(
                StackedShiftBlocks(
                    epsilon=epsilon,
                    mlp_dropout_rate=mlp_dropout_rate,
                    num_shift_blocks=num_shift_blocks,
                    stochastic_depth_rate=stochastic_depth_rate,
                    is_merge=is_merge,
                    num_div=num_div,
                    shift_pixel=shift_pixel,
                    mlp_expand_ratio=mlp_expand_ratio,
                )
            )
        self.global_avg_pool = layers.GlobalAveragePooling2D()

        self.classifier = layers.Dense(config.num_classes)

    def get_config(self):
        config = super().get_config()
        config.update(
            {
                "data_augmentation": self.data_augmentation,
                "patch_projection": self.patch_projection,
                "stages": self.stages,
                "global_avg_pool": self.global_avg_pool,
                "classifier": self.classifier,
            }
        )
        return config

    def _calculate_loss(self, data, training=False):
        (images, labels) = data

        # Augment the images
        augmented_images = self.data_augmentation(images, training=training)

        # Create patches and project the pathces.
        projected_patches = self.patch_projection(augmented_images)

        # Pass through the stages
        x = projected_patches
        for stage in self.stages:
            x = stage(x, training=training)

        # Get the logits.
        x = self.global_avg_pool(x)
        logits = self.classifier(x)

        # Calculate the loss and return it.
        total_loss = self.compiled_loss(labels, logits)
        return total_loss, labels, logits

    def train_step(self, inputs):
        with tf.GradientTape() as tape:
            total_loss, labels, logits = self._calculate_loss(
                data=inputs, training=True
            )

        # Apply gradients.
        train_vars = [
            self.data_augmentation.trainable_variables,
            self.patch_projection.trainable_variables,
            self.global_avg_pool.trainable_variables,
            self.classifier.trainable_variables,
        ]
        train_vars = train_vars + [stage.trainable_variables for stage in self.stages]

        # Optimize the gradients.
        grads = tape.gradient(total_loss, train_vars)
        trainable_variable_list = []
        for (grad, var) in zip(grads, train_vars):
            for g, v in zip(grad, var):
                trainable_variable_list.append((g, v))
        self.optimizer.apply_gradients(trainable_variable_list)

        # Update the metrics
        self.compiled_metrics.update_state(labels, logits)
        return {m.name: m.result() for m in self.metrics}

    def test_step(self, data):
        _, labels, logits = self._calculate_loss(data=data, training=False)

        # Update the metrics
        self.compiled_metrics.update_state(labels, logits)
        return {m.name: m.result() for m in self.metrics}

    def call(self, images):
        augmented_images = self.data_augmentation(images)
        x = self.patch_projection(augmented_images)
        for stage in self.stages:
            x = stage(x, training=False)
        x = self.global_avg_pool(x)
        logits = self.classifier(x)
        return logits

model = ShiftViTModel(
    data_augmentation=get_augmentation_model(),
    projected_dim=config.projected_dim,
    patch_size=config.patch_size,
    num_shift_blocks_per_stages=config.num_shift_blocks_per_stages,
    epsilon=config.epsilon,
    mlp_dropout_rate=config.mlp_dropout_rate,
    stochastic_depth_rate=config.stochastic_depth_rate,
    num_div=config.num_div,
    shift_pixel=config.shift_pixel,
    mlp_expand_ratio=config.mlp_expand_ratio,
)

# 一些代码来源于：
# https://www.kaggle.com/ashusma/training-rfcx-tensorflow-tpu-effnet-b2.
class WarmUpCosine(keras.optimizers.schedules.LearningRateSchedule):
    """一个使用预热余弦衰减调度的学习率调度类。"""

    def __init__(self, lr_start, lr_max, warmup_steps, total_steps):
        """
        参数：
            lr_start: 初始学习率
            lr_max: 在预热步骤中学习率应增加到的最大学习率
            warmup_steps: 模型热身的步骤数
            total_steps: 模型训练的总步骤数
        """
        super().__init__()
        self.lr_start = lr_start
        self.lr_max = lr_max
        self.warmup_steps = warmup_steps
        self.total_steps = total_steps
        self.pi = tf.constant(np.pi)

    def __call__(self, step):
        # 检查总步骤数是否大于预热步骤。如果不是，则抛出值错误。
        if self.total_steps < self.warmup_steps:
            raise ValueError(
                f"总步骤数 {self.total_steps} 必须"
                + f"大于或等于预热步骤数 {self.warmup_steps}。"
            )

        # `cos_annealed_lr` 是一个图形，从初始步骤增加到 1，直到预热步骤。之后，该图形在最后步骤标记处衰减到 -1。
        cos_annealed_lr = tf.cos(
            self.pi
            * (tf.cast(step, tf.float32) - self.warmup_steps)
            / tf.cast(self.total_steps - self.warmup_steps, tf.float32)
        )

        # 将 `cos_annealed_lr` 图形的均值移到 1。现在该图形从 0 变为 2。用 0.5 进行归一化，使其现在从 0 变为 1。通过归一化图形，我们用 `lr_max` 进行缩放，使其从 0 变为 `lr_max`
        learning_rate = 0.5 * self.lr_max * (1 + cos_annealed_lr)

        # 检查 warmup_steps 是否大于 0。
        if self.warmup_steps > 0:
            # 检查 lr_max 是否大于 lr_start。如果不是，则抛出值错误。
            if self.lr_max < self.lr_start:
                raise ValueError(
                    f"lr_start {self.lr_start} 必须小于或"
                    + f"等于 lr_max {self.lr_max}。"
                )

            # 计算学习率在预热调度中应增加的斜率。斜率的公式是 m = ((b-a)/steps)
            slope = (self.lr_max - self.lr_start) / self.warmup_steps

            # 使用直线的公式 (y = mx+c) 构建预热调度
            warmup_rate = slope * tf.cast(step, tf.float32) + self.lr_start

            # 当当前步骤小于预热步骤时，获取直线图。当当前步骤大于预热步骤时，获取缩放后的余弦图。
            learning_rate = tf.where(
                step < self.warmup_steps, warmup_rate, learning_rate
            )

        # 当当前步骤大于总步骤时，返回 0，否则返回计算出的图形。
        return tf.where(
            step > self.total_steps, 0.0, learning_rate, name="learning_rate"
        )

    def get_config(self):
        config = {
            "lr_start": self.lr_start,
            "lr_max": self.lr_max,
            "total_steps": self.total_steps,
            "warmup_steps": self.warmup_steps,
        }
        return config

# pass sample data to the model so that input shape is available at the time of
# saving the model
sample_ds, _ = next(iter(train_ds))
model(sample_ds, training=False)

# Get the total number of steps for training.
total_steps = int((len(x_train) / config.batch_size) * config.epochs)

# Calculate the number of steps for warmup.
warmup_epoch_percentage = 0.15
warmup_steps = int(total_steps * warmup_epoch_percentage)

# Initialize the warmupcosine schedule.
scheduled_lrs = WarmUpCosine(
    lr_start=1e-5,
    lr_max=1e-3,
    warmup_steps=warmup_steps,
    total_steps=total_steps,
)

# Get the optimizer.
optimizer = tfa.optimizers.AdamW(
    learning_rate=scheduled_lrs, weight_decay=config.weight_decay
)

# Compile and pretrain the model.
model.compile(
    optimizer=optimizer,
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=[
        keras.metrics.SparseCategoricalAccuracy(name="accuracy"),
        keras.metrics.SparseTopKCategoricalAccuracy(5, name="top-5-accuracy"),
    ],
)

# Train the model
history = model.fit(
    train_ds,
    epochs=config.epochs,
    validation_data=val_ds,
    callbacks=[
        keras.callbacks.EarlyStopping(
            monitor="val_accuracy",
            patience=5,
            mode="auto",
        )
    ],
)

# Evaluate the model with the test dataset.
print("TESTING")
loss, acc_top1, acc_top5 = model.evaluate(test_ds)
print(f"Loss: {loss:0.2f}")
print(f"Top 1 test accuracy: {acc_top1*100:0.2f}%")
print(f"Top 5 test accuracy: {acc_top5*100:0.2f}%")

model.save("ShiftViT")

!wget -q 'https://tinyurl.com/2p9483sw' -O inference_set.zip
!unzip -q inference_set.zip

# 自定义对象在模型保存时未被包含。
# 在加载时，需要传入这些对象以重建模型
saved_model = tf.keras.models.load_model(
    "ShiftViT",
    custom_objects={"WarmUpCosine": WarmUpCosine, "AdamW": tfa.optimizers.AdamW},
)

def process_image(img_path):
    # 从字符串路径读取图像文件
    img = tf.io.read_file(img_path)

    # 将jpeg解码为uint8张量
    img = tf.io.decode_jpeg(img, channels=3)

    # 调整图像大小以匹配模型接受的输入大小
    # 使用`method`为`nearest`以保留传递给`resize()`的输入的dtype
    img = tf.image.resize(
        img, [config.input_shape[0], config.input_shape[1]], method="nearest"
    )
    return img


def create_tf_dataset(image_dir):
    data_dir = pathlib.Path(image_dir)

    # 使用图像目录创建tf.data数据集
    predict_ds = tf.data.Dataset.list_files(str(data_dir / "*.jpg"), shuffle=False)

    # 使用map将字符串路径转换为uint8图像张量
    # 设置`num_parallel_calls`有助于并行处理多张图像
    predict_ds = predict_ds.map(process_image, num_parallel_calls=AUTO)

    # 创建Prefetch数据集以获得更好的延迟和吞吐量
    predict_ds = predict_ds.batch(config.tf_ds_batch_size).prefetch(AUTO)
    return predict_ds


def predict(predict_ds):
    # ShiftViT模型返回logits（非标准化预测）
    logits = saved_model.predict(predict_ds)

    # 通过调用softmax()对预测进行归一化
    probabilities = tf.nn.softmax(logits)
    return probabilities


def get_predicted_class(probabilities):
    pred_label = np.argmax(probabilities)
    predicted_class = config.label_map[pred_label]
    return predicted_class


def get_confidence_scores(probabilities):
    # 获取按降序排列的概率分数的索引
    labels = np.argsort(probabilities)[::-1]
    confidences = {
        config.label_map[label]: np.round((probabilities[label]) * 100, 2)
        for label in labels
    }
    return confidences

img_dir = "inference_set"
predict_ds = create_tf_dataset(img_dir)
probabilities = predict(predict_ds)
print(f"probabilities: {probabilities[0]}")
confidences = get_confidence_scores(probabilities[0])
print(confidences)

plt.figure(figsize=(10, 10))
for images in predict_ds:
    for i in range(min(6, probabilities.shape[0])):
        ax = plt.subplot(3, 3, i + 1)
        plt.imshow(images[i].numpy().astype("uint8"))
        predicted_class = get_predicted_class(probabilities[i])
        plt.title(predicted_class)
        plt.axis("off")