Keras实现Vision Transformer VIT模型示例详解_python_

#--------------------------------------------------------------------------------------------------------------------# # classtoken部分是transformer的分类特征。用于堆叠到序列化后的图片特征中，作为一个单位的序列特征进行特征提取。 # # 在利用步长为16x16的卷积将输入图片划分成14x14的部分后，将14x14部分的特征平铺，一幅图片会存在序列长度为196的特征。 # 此时生成一个classtoken，将classtoken堆叠到序列长度为196的特征上，获得一个序列长度为197的特征。 # 在特征提取的过程中，classtoken会与图片特征进行特征的交互。最终分类时，我们取出classtoken的特征，利用全连接分类。 #--------------------------------------------------------------------------------------------------------------------# class ClassToken(Layer): def __init__(self, cls_initializer='zeros', cls_regularizer=None, cls_constraint=None, **kwargs): super(ClassToken, self).__init__(**kwargs) self.cls_initializer = keras.initializers.get(cls_initializer) self.cls_regularizer = keras.regularizers.get(cls_regularizer) self.cls_constraint = keras.constraints.get(cls_constraint) def get_config(self): config = { 'cls_initializer': keras.initializers.serialize(self.cls_initializer), 'cls_regularizer': keras.regularizers.serialize(self.cls_regularizer), 'cls_constraint': keras.constraints.serialize(self.cls_constraint), } base_config = super(ClassToken, self).get_config() return dict(list(base_config.items()) + list(config.items())) def compute_output_shape(self, input_shape): return (input_shape[0], input_shape[1] + 1, input_shape[2]) def build(self, input_shape): self.num_features = input_shape[-1] self.cls = self.add_weight( shape = (1, 1, self.num_features), initializer = self.cls_initializer, regularizer = self.cls_regularizer, constraint = self.cls_constraint, name = 'cls', ) super(ClassToken, self).build(input_shape) def call(self, inputs): batch_size = tf.shape(inputs)[0] cls_broadcasted = tf.cast(tf.broadcast_to(self.cls, [batch_size, 1, self.num_features]), dtype = inputs.dtype) return tf.concat([cls_broadcasted, inputs], 1) #--------------------------------------------------------------------------------------------------------------------# # 为网络提取到的特征添加上位置信息。 # 以输入图片为224, 224, 3为例，我们获得的序列化后的图片特征为196, 768。加上classtoken后就是197, 768 # 此时生成的pos_Embedding的shape也为197, 768，代表每一个特征的位置信息。 #--------------------------------------------------------------------------------------------------------------------# class AddPositionEmbs(Layer): def __init__(self, image_shape, patch_size, pe_initializer='zeros', pe_regularizer=None, pe_constraint=None, **kwargs): super(AddPositionEmbs, self).__init__(**kwargs) self.image_shape = image_shape self.patch_size = patch_size self.pe_initializer = keras.initializers.get(pe_initializer) self.pe_regularizer = keras.regularizers.get(pe_regularizer) self.pe_constraint = keras.constraints.get(pe_constraint) def get_config(self): config = { 'pe_initializer': keras.initializers.serialize(self.pe_initializer), 'pe_regularizer': keras.regularizers.serialize(self.pe_regularizer), 'pe_constraint': keras.constraints.serialize(self.pe_constraint), } base_config = super(AddPositionEmbs, self).get_config() return dict(list(base_config.items()) + list(config.items())) def compute_output_shape(self, input_shape): return input_shape def build(self, input_shape): assert (len(input_shape) == 3), f"Number of dimensions should be 3, got {len(input_shape)}" length = (224 // self.patch_size) * (224 // self.patch_size) + 1 self.pe = self.add_weight( # shape = [1, input_shape[1], input_shape[2]], shape = [1, length, input_shape[2]], initializer = self.pe_initializer, regularizer = self.pe_regularizer, constraint = self.pe_constraint, name = 'pos_embedding', ) super(AddPositionEmbs, self).build(input_shape) def call(self, inputs): num_features = tf.shape(inputs)[2] cls_token_pe = self.pe[:, 0:1, :] img_token_pe = self.pe[:, 1: , :] img_token_pe = tf.reshape(img_token_pe, [1, (224 // self.patch_size), (224 // self.patch_size), num_features]) img_token_pe = tf.image.resize_bicubic(img_token_pe, (self.image_shape[0] // self.patch_size, self.image_shape[1] // self.patch_size), align_corners=False) img_token_pe = tf.reshape(img_token_pe, [1, -1, num_features]) pe = tf.concat([cls_token_pe, img_token_pe], axis = 1) return inputs + tf.cast(pe, dtype=inputs.dtype) def VisionTransformer(input_shape = [224, 224], patch_size = 16, num_layers = 12, num_features = 768, num_heads = 12, mlp_dim = 3072, classes = 1000, dropout = 0.1): #-----------------------------------------------# # 224, 224, 3 #-----------------------------------------------# inputs = Input(shape = (input_shape[0], input_shape[1], 3)) #-----------------------------------------------# # 224, 224, 3 -> 14, 14, 768 #-----------------------------------------------# x = Conv2D(num_features, patch_size, strides = patch_size, padding = "valid", name = "patch_embed.proj")(inputs) #-----------------------------------------------# # 14, 14, 768 -> 196, 768 #-----------------------------------------------# x = Reshape(((input_shape[0] // patch_size) * (input_shape[1] // patch_size), num_features))(x) #-----------------------------------------------# # 196, 768 -> 197, 768 #-----------------------------------------------# x = ClassToken(name="cls_token")(x) #-----------------------------------------------# # 197, 768 -> 197, 768 #-----------------------------------------------# x = AddPositionEmbs(input_shape, patch_size, name="pos_embed")(x)