论文阅读笔记——ViT

ViT Block

• Patch Embedding（将视觉问题转化为 seq2seq）：假设输入图片为 224 * 224，将图片分为固定大小 patch=16 * 16，则每张图片生成 224 * 224 / 16 * 16 = 196个patch，即输入序列长度为196（dim=16 * 16 * 3= 768），那么线性投影层维度为 768 * N（N=768），故通过线性投影层厚为 196 * 768（196个token，每个token维度768），为模仿 BERT，加上特殊字符 cls，故最终为 197 * 768
• Positional Embedding：位置编码与 patch 为 sum关系，故仍未 197 * 768
• LN/Multi-Head Attention/LN：多头自注意力时，将输入映射到qkv，如果有12个头，则qkv维度为197 * 64（768 / 12 = 64），最后将12组的qkv拼接起来，仍然为 197 * 768
• MLP：将 197 * 768 放大为 197 * 3072，再缩小为 197 * 768
  $$\begin{array}{rlrlr}TheMLPcontains& twolayerswithaGELUnon-linearity.& & & \\ & {\mathbf{z}}_0=[{\mathbf{x}}_{\mathrm{class}};{\mathbf{x}}_p^1\mathbf{E};{\mathbf{x}}_p^2\mathbf{E};\cdots ;{\mathbf{x}}_p^N\mathbf{E}]+{\mathbf{E}}_{pos},\mathbf{I}& \mathbf{E}\in {\mathbb{R}}^{(P^2\cdot C)\times D},{\mathbf{E}}_{pos}\in {\mathbb{R}}^{(N+1)\times D}& & (1)\\ & {{\mathrm{z}}^{\prime }}_{\ell }=\mathrm{MSA}(\mathrm{LN}({\mathbf{z}}_{\ell -1}))+{\mathbf{z}}_{\ell -}& \ell =1\dots L& & (2)\\ & {\mathbf{z}}_{\ell }=\mathrm{MLP}(\mathrm{LN}({{\mathbf{z}}^{\prime }}_{\ell }))+{{\mathbf{z}}^{\prime }}_{\ell },& \ell =1\dots L& & (3)\\ & \mathbf{y}=\mathrm{LN}({\mathbf{z}}_L^0)& (4)& & \end{array}$$

相比于 CNN 的优势

• 捕捉全局特征
• 适应不同输入尺寸

当拥有足够多的数据进行训练时，ViT 的表现会超过 CNN；当数据量不够时，ViT 的表现弱于 CNN，因为 transformer 相比于 CNN 缺少局部性和平移不变性两种假设

代码实例

class AddPositionEmbs(nn.Module):
	# posemb_init: 一个可调用对象，用于初始化位置嵌入矩阵。它接受PRNGKey（随机数生成器密钥）、Shape（形状）和Dtype（数据类型）作为参数，并返回一个Array（数组）
    posemb_init: Callable[[PRNGKey, Shape, Dtype], Array]
    param_dtype : Dtype = jnp.float32 # 参数的数据类型，默认为jnp.float32
    
    @nn.compact
    def __call__(self, inputs):
        assert inputs.ndim == 3, f'Expected (batch, seq, hidden) got {inputs.shape}' # inputs （batch_size, seq_len, emb_dim）
        pos_emb_shape = (1, inputs.shape[1], inputs.shape[2]) # 位置嵌入 （1, seq_len, emb_dim）
        pe = self.param('pos_embedding', self.posemb_init, pos_emb_shape, self.param_dtype)
        return inputs + pe

class MlpBlock(nn.Module):
    """Transformer MLP / feed-forward block."""

    mlp_dim: int
    dtype: Dtype = jnp.float32
    param_dtype: Dtype = jnp.float32
    out_dim: Optional[int] = None
    dropout_rate: float = 0.1
    kernel_init: Callable[[PRNGKey, Shape, Dtype], Array] = nn.initializers.xavier_uniform()
    bias_init: Callable[[PRNGKey, Shape, Dtype], Array] = nn.initializers.normal(stddev=1e-6)

    @nn.compact
    def __call__(self, inputs, *, deterministic):
    	# 全连接 + gelu + Dropout + 全连接 + Dropout
        actual_out_dim = inputs.shape[-1] if self.out_dim is None else self.out_dim
        x = nn.Dense(features=self.mlp_dim, dtype=self.dtype, param_dtype=self.param_dtype,
                     kernel_init=self.kernel_init, bias_init=self.bias_init)(inputs)
        x = nn.gelu(x)
        x = nn.Dropout(rate=self.dropout_rate)(x, deterministic=deterministic)
        output = nn.Dense(features=actual_out_dim, dtype=self.dtype, param_dtype=self.param_dtype,
                          kernel_init=self.kernel_init, bias_init=self.bias_init)(x)
        output = nn.Dropout(rate=self.dropout_rate)(output, deterministic=deterministic)
        return output

class Encoder1DBlock(nn.Module):
    """Transformer encoder layer."""

    mlp_dim: int
    num_heads: int
    dtype: Dtype = jnp.float32
    dropout_rate: float = 0.1
    attention_dropout_rate: float = 0.1

    @nn.compact
    def __call__(self, inputs, *, deterministic):
    	# 层归一化 + 多头自注意力 + Dropout + 残差链接 + 层归一化 + MLP + 残差链接
        assert inputs.ndim == 3, f'Expected (batch, seq, hidden) got {inputs.shape}'
        x = nn.LayerNorm(dtype=self.dtype)(inputs)
        # 多头自注意力
        x = nn.MultiHeadDotProductAttention(dtype=self.dtype, kernel_init=nn.initializers.xavier_uniform(), # 权重初始化方法
                                            broadcast_dropout=False,  # 是否广播 Dropout
                                            deterministic=deterministic, # 是否启用 Dropout
                                            dropout_rate=self.attention_dropout_rate, num_heads=self.num_heads)(x, x)
        x = nn.Dropout(rate=self.dropout_rate)(x, deterministic=deterministic)
        x = x + inputs

        y = nn.LayerNorm(dtype=self.dtype)(x)
        y = MlpBlock(mlp_dim=self.mlp_dim, dtype=self.dtype, dropout_rate=self.dropout_rate)(y, deterministic=deterministic)
        return x + y

# Transformer Encoder
class Encoder(nn.Module):
    """Transformer Model Encoder for sequence to sequence translation."""

    num_layers: int
    mlp_dim: int
    num_heads: int
    dropout_rate: float = 0.1
    attention_dropout_rate: float = 0.1
    add_position_embedding: bool = True

    @nn.compact
    def __call__(self, x, *, train):
        assert x.ndim == 3, f'Expected (batch, len, emb) got {x.shape}
        # 是否添加位置嵌入'
        if self.add_position_embedding:
            x = AddPositionEmbs(posemb_init=nn.initializers.normal(stddev=0.02))(x)
            x = nn.Dropout(rate=self.dropout_rate)(x, deterministic=not train)

        for lyr in range(self.num_layers):
            x = Encoder1DBlock(mlp_dim=self.mlp_dim, dropout_rate=self.dropout_rate,
                               attention_dropout_rate=self.attention_dropout_rate,
                               num_heads=self.num_heads)(x, deterministic=not train)
        encoded = nn.LayerNorm(name='encoder_norm')(x)
        return encoded

将输入图像通过可选的 ResNet 预处理阶段提取特征，然后将图像分割成 patch 并嵌入到 Transformer 可处理的序列中。通过 Transformer 编码器对序列进行编码后，根据分类器的类型（如 token、gap 等）提取特征表示，最后通过全连接层输出分类结果。该模型支持添加类别 token、全局平均池化等操作，并可选择是否在 Transformer 前使用 ResNet 进行特征提取。

class VisionTransformer(nn.Module):
  """VisionTransformer."""

  num_classes: int
  patches: Any
  transformer: Any
  hidden_size: int
  resnet: Optional[Any] = None
  representation_size: Optional[int] = None
  classifier: str = 'token'
  head_bias_init: float = 0.
  encoder: Type[nn.Module] = Encoder
  model_name: Optional[str] = None

  @nn.compact
  def __call__(self, inputs, *, train):

    x = inputs
    # (Possibly partial) ResNet root.
    if self.resnet is not None:
      width = int(64 * self.resnet.width_factor)

      # Root block.
      x = models_resnet.StdConv(
          features=width,
          kernel_size=(7, 7),
          strides=(2, 2),
          use_bias=False,
          name='conv_root')(
              x)
      x = nn.GroupNorm(name='gn_root')(x)
      x = nn.relu(x)
      x = nn.max_pool(x, window_shape=(3, 3), strides=(2, 2), padding='SAME')

      # ResNet stages.
      if self.resnet.num_layers:
        x = models_resnet.ResNetStage(
            block_size=self.resnet.num_layers[0],
            nout=width,
            first_stride=(1, 1),
            name='block1')(
                x)
        for i, block_size in enumerate(self.resnet.num_layers[1:], 1):
          x = models_resnet.ResNetStage(
              block_size=block_size,
              nout=width * 2**i,
              first_stride=(2, 2),
              name=f'block{i + 1}')(
                  x)

    n, h, w, c = x.shape

    # We can merge s2d+emb into a single conv; it's the same.
    x = nn.Conv(
        features=self.hidden_size,
        kernel_size=self.patches.size,
        strides=self.patches.size,
        padding='VALID',
        name='embedding')(
            x)

    # Here, x is a grid of embeddings.

    # (Possibly partial) Transformer.
    if self.transformer is not None:
      n, h, w, c = x.shape
      x = jnp.reshape(x, [n, h * w, c])

      # If we want to add a class token, add it here.
      if self.classifier in ['token', 'token_unpooled']:
        cls = self.param('cls', nn.initializers.zeros, (1, 1, c))
        cls = jnp.tile(cls, [n, 1, 1])
        x = jnp.concatenate([cls, x], axis=1)

      x = self.encoder(name='Transformer', **self.transformer)(x, train=train)

    if self.classifier == 'token':
      x = x[:, 0]
    elif self.classifier == 'gap':
      x = jnp.mean(x, axis=list(range(1, x.ndim - 1)))  # (1,) or (1,2)
    elif self.classifier in ['unpooled', 'token_unpooled']:
      pass
    else:
      raise ValueError(f'Invalid classifier={self.classifier}')

    if self.representation_size is not None:
      x = nn.Dense(features=self.representation_size, name='pre_logits')(x)
      x = nn.tanh(x)
    else:
      x = IdentityLayer(name='pre_logits')(x)

    if self.num_classes:
      x = nn.Dense(
          features=self.num_classes,
          name='head',
          kernel_init=nn.initializers.zeros,
          bias_init=nn.initializers.constant(self.head_bias_init))(x)
    return x