目的

基于上一篇MINICPM-V2_6图像得到embedding-代码解读将图片patch找到对应的embedding（包括位置embedding和像素embedding），embedding经过多层attention后会得到vision_embedding，vision_embedding的长度对应的是patch的个数，这个长度是不固定的，有长有短，那要怎么做到长度统一呢？
这篇从vision_embedding入手，了解如何解决patch长度不统一的问题

Perceiver Resampler

Flamingo

Perceiver Resampler解决的事将变长的patch信息转化为固定大小长度的特征，否则过长的patch会大大加大后续LLM的计算负担(毕竟每个token都占用计算)。
Perceiver Resampler采用一个可学习的queries作为交叉注意力中的Q，而将patch进行特征提取后的表示x_f，和Q拼接起来作为交叉注意力中的K和V，通过这种方法将变长的patch特征就规整为了固定大小的特征，方便了后续的处理。

代码

基本变量

import torch
from torch import nn
from torch import Tensor
from functools import partial
import numpy as np

num_queries = 64# query数量
embed_dim = 3584# 向量维度
num_heads = 3584//128# attention头数28
adaptive = True# 
max_size = (70,70)# 最大尺寸
kv_dim = 1152# kv向量维度

query = nn.Parameter(torch.zeros(num_queries, embed_dim))# 64，3584
kv_proj = nn.Linear(kv_dim, embed_dim, bias=False)# 1152,3584

norm_layer=partial(nn.LayerNorm, eps=1e-6)
ln_q = norm_layer(embed_dim)# LayerNorm((3584,), eps=1e-06, elementwise_affine=True)
ln_kv = norm_layer(embed_dim)# LayerNorm((3584,), eps=1e-06, elementwise_affine=True)

ln_post = norm_layer(embed_dim)# LayerNorm((3584,), eps=1e-06, elementwise_affine=True)
proj = nn.Parameter((embed_dim ** -0.5) * torch.randn(embed_dim, embed_dim))# 3584,3584

这里定义了几个LN，都是在attention前后使用的

函数定义

既然是attention，那其中必然有位置embedding，这里使用的是ROPE，只是因为是2D，所以这里也要处理一下得到2D的位置embedding

def get_2d_sincos_pos_embed(embed_dim, image_size):
    """
    输入：
        embed_dim: 向量维度
        image_size:(H,W)
    输出：
        pos_embed：(H, W, embed_dim)
    demo:
        embed_dim = 8
        image_size = (3,5)
        pos_embed = get_2d_sincos_pos_embed(embed_dim, image_size)# 3,5,8
    """
    if isinstance(image_size, int):
        grid_h_size, grid_w_size = image_size, image_size
    else:
        grid_h_size, grid_w_size = image_size[0], image_size[1]# 70,70
    grid_h = np.arange(grid_h_size, dtype=np.float32)# 0,1,2,..,69
    grid_w = np.arange(grid_w_size, dtype=np.float32)# 0,1,2,..,69
    grid = np.meshgrid(grid_w, grid_h)  # 生成网格，但是这里是w在前；torch.meshgrid是h在前
    # [[ 0.,  1.,  2., ..., 67., 68., 69.], [ 0.,  1.,  2., ..., 67., 68., 69.],
    # [[ 0.,  0.,  0., ...,  0.,  0.,  0.], [ 1.,  1.,  1., ...,  1.,  1.,  1.],
    grid = np.stack(grid, axis=0)# 在第0维拼接
    # [[[ 0.,  1.,  2., ..., 67., 68., 69.], [ 0.,  1.,  2., ..., 67., 68., 69.],
    # [[ 0.,  0.,  0., ...,  0.,  0.,  0.], [ 1.,  1.,  1., ...,  1.,  1.,  1.], ]]   
    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)# 70,70,3584
    return pos_embed


def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
    """
    输入：
        embed_dim: 向量维度
        grid:行位置和列位置[(H,W),(H,W)]
    输出：
        emb：(H, W, embed_dim)
    demo:
        embed_dim = 8
        grid = [[np.arange(70, dtype=np.float32)],[np.arange(70, dtype=np.float32)]]
        emb = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)# 1,70,8
    """
    assert embed_dim % 2 == 0
    # use half of dimensions to encode grid_h
    emb_h = get_1d_sincos_pos_embed_from_grid_new(embed_dim // 2, grid[0])  # (H, W, D/2) H维度是一样的 W维度是不一样的 对应的是左半部分
    emb_w = get_1d_sincos_pos_embed_from_grid_new(embed_dim // 2, grid[1])  # (H, W, D/2) W维度是一样的 H维度是不一样的 对应的是右半部分
    emb = np.concatenate([emb_h, emb_w], axis=-1)  # (H, W, D)
    return emb


def get_1d_sincos_pos_embed_from_grid_new(embed_dim, pos):
    """
    输入：
        embed_dim: 向量维度
        pos: 位置 (H, W)
    输出:
        emb：得到POS对应的ROPE位置向量 (H, W, D)
    demo:
    embed_dim = 8
    pos = [np.arange(70, dtype=np.float32)]# [[0,1,2,..,69]]
    emb = get_1d_sincos_pos_embed_from_grid_new(embed_dim, pos)# 1,70,8
    """
    assert embed_dim % 2 == 0
    omega = np.arange(embed_dim // 2, dtype=np.float32)# 0,1,2,3
    omega /= embed_dim / 2.# [0,1,2,3]/4=[0,1/4,2/4,3/4]
    omega = 1. / 10000 ** omega  # (D/2,) ROPE位置编码中的\theta = 1/10000 **(2i/d)
    out = np.einsum('hw,d->hwd', pos, omega)  # (H, W, D/2), outer product m*\theta
    emb_sin = np.sin(out)  # (H, W, D/2) sin(m*\theta)
    emb_cos = np.cos(out)  # (H, W, D/2) cos(m*\theta)
    emb = np.concatenate([emb_sin, emb_cos], axis=-1)  # (H, W, D)
    return emb

embed_dim = 8
max_size = (3,5)
pos_embed = torch.from_numpy(get_2d_sincos_pos_embed(embed_dim, max_size)).float()

tgt_sizes = (2,3)
tgt_h, tgt_w = tgt_sizes
a = pos_embed[:tgt_h, :tgt_w, :].reshape((tgt_h * tgt_w, -1))# tgt_h * tgt_w, 8

我一直想象不出来这里得到的pos_embed和a是什么样子的，简单举个例子，看着清楚明白
h=3，w=5，embed_dim=8，最后形成的pos_embed的维度是(3,5,8)，相同颜色的是一样的值

当一个图片patch对应的tgt_sizes=(2,3)从pos_embed中找到的位置embedding是什么样子的呢？
将h对应的坐标当作i，将w对应的坐标当作j，从pos_embed找到pos_embed[:h,:w,:]就是下面这个样子的了

def MultiheadAttention(embed_dim, num_heads, query ,key ,value, key_padding_mask):
    """
    输入：
        embed_dim: 向量维度
        num_heads: 头数
    输出:
        output:输出向量 (query_num, batch_size, embed_dim)
    demo:
    这里和attention是一样的，后面找个机会写
    """    
    output = torch.randn(64, 3, embed_dim)
    return output

函数调用

x = torch.randn(3,1036,1152)# 上一步得到的图片patch对应的embedding
tgt_sizes = torch.tensor([[28, 37],
        [39, 26],
        [39, 26]])# 上一步得到的图片patch对应的尺寸
assert x.shape[0] == tgt_sizes.shape[0]
bs = x.shape[0]

embed_dim = 3584# 向量维度
max_size = (70,70)# 最大尺寸
pos_embed = torch.from_numpy(get_2d_sincos_pos_embed(embed_dim, max_size)).float()# 70,70,3584

patch_len = tgt_sizes[:, 0] * tgt_sizes[:, 1]# [1036, 1014, 1014]
max_patch_len = torch.max(patch_len)# 1036
key_padding_mask = torch.zeros((bs, max_patch_len))

position_embed = []# 通过从pos_embed中根据图片patch大小得到对应的位置向量，且要更新key_padding_mask，后面的attention要用的
for i in range(bs):
    tgt_h, tgt_w = tgt_sizes[i]
    position_embed.append(pos_embed[:tgt_h, :tgt_w, :].reshape((tgt_h * tgt_w, -1)))  # [[1036, 3584],[1014, 3584],[1014, 3584]]
    key_padding_mask[i, patch_len[i]:] = True

position_embed = torch.nn.utils.rnn.pad_sequence(
    position_embed, batch_first=True, padding_value=0.0).permute(1, 0, 2)  # BLD => L * B * D 1036*3*3584

x = kv_proj(x)  # 将x映射得到kv的过程 [3, 1036, 3584]
x = ln_kv(x).permute(1, 0, 2)  # 将x经过ln [1036, 3, 3584]

q = ln_q(query)  # 将query经过ln 64*3584

def repeat(query, N: int):
    return query.unsqueeze(1).repeat(1, N, 1)


out = MultiheadAttention(embed_dim, num_heads,
    repeat(q, bs),  # Q * B * D query
    x + position_embed,  # L * B * D +  L * B * D key
    x,# L * B * D  value
    key_padding_mask=key_padding_mask)
#  out: Q * B * D 在这一步完成了映射
# attn就是正常的attention操作了
# key、value的值和原始论文flamingo中有些不一样，论文中是将q和x拼接在一起作为key、value进行attention，这里只使用了x
# 位置编码的使用和论文也是不一样的，论文中是将位置编码直接加在了原始输入上，所以value、key都是有位置编码的
# 这里是只将位置编码加在了key上


x = out.permute(1, 0, 2)  # B * Q * D

x = ln_post(x)# ln
x = x @ proj# 映射
return x

额外说几句

# 模型MINICPM-V2_6中的resampler部分结构
  (resampler): Resampler(
    (kv_proj): Linear(in_features=1152, out_features=3584, bias=False)
    (attn): MultiheadAttention(
      (out_proj): Linear(in_features=3584, out_features=3584, bias=True)
    )
    (ln_q): LayerNorm((3584,), eps=1e-06, elementwise_affine=True)
    (ln_kv): LayerNorm((3584,), eps=1e-06, elementwise_affine=True)
    (ln_post): LayerNorm((3584,), eps=1e-06, elementwise_affine=True)
  )

这里主要是为了将变长的vision embedding转化为固定长度，代码结构也比较简单
代码比较简单，能想通2D position embeeding就好。
挖了一个坑，要看看attention的代码了