ViT模型的出现标志着Transformer架构在计算机视觉中的成功应用

import tensorflow as tf
from tensorflow.keras import layers, models, initializers
import numpy as np
 
 
class PatchEmbedding(layers.Layer):
    def __init__(self, patch_size, num_patches, embed_dim):
        super(PatchEmbedding, self).__init__()
        self.num_patches = num_patches
        self.proj = layers.Dense(embed_dim)
        self.cls_token = self.add_weight("cls_token", shape=[1, 1, embed_dim], initializer=initializers.Zeros())
        self.pos_embed = self.add_weight("pos_embed", shape=[1, num_patches + 1, embed_dim], initializer=initializers.Zeros())
 
    def call(self, x):
        batch_size, height, width, channels = x.shape
        patch_size_h, patch_size_w = height // self.num_patches, width // self.num_patches
        x = tf.image.extract_patches(x, sizes=[1, patch_size_h, patch_size_w, 1], strides=[1, patch_size_h, patch_size_w, 1], rates=[1, 1, 1, 1], padding='VALID')
        x = tf.reshape(x, [batch_size, -1, patch_size_h * patch_size_w * channels])
        x = self.proj(x)
 
        cls_tokens = tf.broadcast_to(self.cls_token, [batch_size, 1, self.proj.units])
        x = tf.concat([cls_tokens, x], axis=1)
        x += self.pos_embed
        return x
 
 
class MultiHeadSelfAttention(layers.Layer):
    def __init__(self, embed_dim, num_heads):
        super(MultiHeadSelfAttention, self).__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.proj_qkv = layers.Dense(3 * embed_dim)
        self.proj_out = layers.Dense(embed_dim)
 
    def call(self, x):
        batch_size, num_patches, embed_dim = x.shape
        qkv = self.proj_qkv(x)
        q, k, v = tf.split(qkv, 3, axis=-1)
        q = self.split_heads(q)
        k = self.split_heads(k)
        v = self.split_heads(v)
        attention_scores = tf.einsum('bhqd,bhkd->bhqk', q, k) / tf.math.sqrt(float(embed_dim))
        attention_weights = tf.nn.softmax(attention_scores, axis=-1)
        attention_output = tf.einsum('bhqk,bhvd->bhqd', attention_weights, v)
        attention_output = self.combine_heads(attention_output)
        return self.proj_out(attention_output)
 
    def split_heads(self, x):
        batch_size, num_patches, embed_dim = x.shape
        depth = embed_dim // self.num_heads
        x = tf.reshape(x, [batch_size, num_patches, self.num_heads, depth])
        return tf.transpose(x, perm=[0, 2, 1, 3])
 
    def combine_heads(self, x):
        batch_size, num_heads, num_patches, depth = x.shape
        x = tf.transpose(x, perm=[0, 2, 1, 3])
        return tf.reshape(x, [batch_size, num_patches, num_heads * depth])
 
 
class TransformerBlock(layers.Layer):
    def __init__(self, embed_dim, num_heads, mlp_dim, dropout_rate):
        super(TransformerBlock, self).__init__()
        self.mha = MultiHeadSelfAttention(embed_dim, num_heads)
        self.mlp = models.Sequential([
            layers.Dense(mlp_dim, activation=tf.nn.gelu),
            layers.Dense(embed_dim)
        ])
        self.layernorm1 = layers.LayerNormalization(epsilon=1e-6)
        self.layernorm2 = layers.LayerNormalization(epsilon=1e-6)
        self.dropout1 = layers.Dropout(dropout_rate)
        self.dropout2 = layers.Dropout(dropout_rate)
 
    def call(self, x, training):
        attn_output = self.mha(self.layernorm1(x))
        attn_output = self.dropout1(attn_output, training=training)
        out1 = x + attn_output
        mlp_output = self.mlp(self.layernorm2(out1))
        mlp_output = self.dropout2(mlp_output, training=training)
        return out1 + mlp_output
 
 
def create_vit_model(input_shape, patch_size, num_layers, num_patches, embed_dim, num_heads, mlp_dim, num_classes, dropout_rate):
    inputs = layers.Input(shape=input_shape)
    x = PatchEmbedding(patch_size, num_patches, embed_dim)(inputs)
 
    for _ in range(num_layers):
        x = TransformerBlock(embed_dim, num_heads, mlp_dim, dropout_rate)(x)
 
    x = layers.LayerNormalization(epsilon=1e-6)(x)
    x = x[:, 0]
    x = layers.Dense(num_classes)(x)
 
    model = models.Model(inputs=inputs, outputs=x)
    return model
 
 
# 超参数
input_shape = (224, 224, 3)
patch_size = 16
num_layers = 12
num_patches = (input_shape[0] // patch_size) * (input_shape[1] // patch_size)
embed_dim = 768
num_heads = 12
mlp_dim = 3072
num_classes = 10
dropout_rate = 0.1
 
vit_model = create_vit_model(input_shape, patch_size, num_layers, num_patches, embed_dim, num_heads, mlp_dim, num_classes, dropout_rate)
vit_model.compile(optimizer='adam', loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), metrics=['accuracy'])
 
vit_model.summary()