图像特征压缩的多种实现

《DeepSeek大模型高性能核心技术与多模态融合开发(人工智能技术丛书)》(王晓华)【摘要 书评 试读】- 京东图书

图像特征token作为图文多模态特征之一,在视觉与语言的联合表示学习中扮演着至关重要的角色。它们能够有效地捕捉图像的局部细节和全局信息,为跨模态检索、视觉问答以及图像描述生成等任务提供了丰富的特征支持。

在图像特征token的研究中,我们不断探索如何更有效地提取、表示和利用这些特征。从最初的底层视觉特征,如边缘、纹理和颜色,到后来的高层语义特征,如物体、场景和动作,图像特征token的表达能力逐渐增强,为模型提供了更丰富和更准确的输入信息。

但是,随着数据规模的扩大和模型复杂度的提升,图像特征token的处理也面临着前所未有的挑战。如何在保证特征表达能力的同时,降低计算成本、提高处理效率,成为了当前研究的重要课题。

12.1.1  Pixel-Shuffle的token压缩

作为图像token压缩的一种高效策略,我们的核心目标在于直接削减token的数量,同时确保压缩后的token集合仍能充分承载原始图像的关键信息,维持信息的完整性和表达力。在这一理念的指引下,Pixel-Shuffle技术应运而生,成为实现这一目标的具体手段。

Pixel-Shuffle,顾名思义,是一种通过“像素重组”来实现数据维度转换的创新方法。其核心思想在于巧妙地利用通道与空间之间的转换关系,通过牺牲一定的空间分辨率来换取通道数的增加。这一转换过程不仅有效地降低了数据的空间维度,还使得每个token能够蕴含更为丰富和紧凑的图像特征,从而在减少token数目的同时,保证了图像信息的容量和表达质量。

具体来说,Pixel-Shuffle技术通过一系列精心设计的操作,将原始图像数据在通道维度上进行扩展,并在空间维度上进行相应的缩减。这种维度变换不仅有助于减少数据的冗余度,还能提升模型对关键特征的敏感度和捕捉能力。经过Pixel-Shuffle处理后的图像token,不仅数量得到了有效控制,而且每个token都蕴含了更为精炼和有价值的图像信息,为后续的模型处理提供了更优质和更高效的输入。其实现如下:

# 定义一个名为pixel_shuffle的函数,它接受一个输入张量x和一个缩放因子scale_factor,默认值为2。
def pixel_shuffle(x, scale_factor=2):
    # 使用einops.rearrange函数将输入张量x从(batch, channel, height, width)格式重新排列为(batch, height, width, channel)格式。
    x = einops.rearrange(x,"b c h w -> b h w c")
    # 获取重新排列后张量的维度信息,分别是批次大小n,宽度w,高度h和通道数c。
    n, w, h, c = x.size()

    # 再次使用einops.rearrange函数将张量按照(scale_factor, scale_factor)的块进行重组,并将通道数扩展为原来的(scale_factor * scale_factor)倍。
    x = einops.rearrange(x,"b h (w s) c -> b h w (c s)",s = scale_factor)
    # 调整张量的维度顺序,使其变为(batch, width, height, channel)格式,并确保数据在内存中是连续的。
    x = x.permute(0, 2, 1, 3).contiguous()
    # 调整张量的形状,使其高度和宽度分别除以缩放因子,通道数乘以缩放因子的平方。
    x = x.view(n, int(h // scale_factor), int(w // scale_factor),int(c * (scale_factor * scale_factor)))

    # 再次调整张量的维度顺序,使其变为(batch, height, width, channel)格式,并确保数据在内存中是连续的。
    x = x.permute(0, 2, 1, 3).contiguous()
    # 最后,使用einops.rearrange函数将张量重新排列回(batch, channel, height, width)格式。
    x = einops.rearrange(x,"b h w c-> b c h w")
    # 返回处理后的张量。
    return x

上面代码定义了一个名为pixel_shuffle的函数,它接受一个输入张量x和一个缩放因子scale_factor,通过重新排列和调整张量的形状,实现了像素洗牌操作,这是一种用于图像上采样的技术,可以增加图像的空间分辨率而不增加计算量。

下面采用Patch_Embedding的方式,来计算图像的token数。我们首先计算没有经过变换前的图像,之后对比经过重排的图像,完整代码如下所示。

import torch
import einops

class PatchEmbedding(torch.nn.Module):
    def __init__(self, image_size, in_channels, patch_size = 14 , embed_dim = 312, dropout=0.):
        super(PatchEmbedding, self).__init__()
        # patch_embed相当于做了一个卷积
        self.patch_embed = torch.nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size, bias=False)
        self.drop = torch.nn.Dropout(dropout)

    def forward(self, x):
        # x[4, 3, 224, 224]
        x = self.patch_embed(x)
        # x [4, 16, 32, 32]
        # x:[n,embed_dim,h',w']
        x = x.flatten(2)  # 将x拉直,h'和w'合并   [n,embed,h'*w']   #x [4, 16, 1024]
        x = x.permute(0, 2, 1)  # [n,h'*w',embed]      #x [4, 1024, 16]

        return x

if __name__ == '__main__':
    image = torch.randn(size=(2,3,224,224))
    image_token = PatchEmbedding(224,3)(image)
    print(image_token.shape)

    image = pixel_shuffle(image)
    image_token = PatchEmbedding(112,12)(image)
    print(image_token.shape)

结果如下:

torch.Size([2, 256, 312])

torch.Size([2, 64, 312])

对比未经过变换前的token化图像,经过转换后的图像在整体维度上缩小了3/4,而只有原有的1/4。其缩减的大小则是由缩放比例scale_factor得到,我们可以根据需要调整不同大小的scale_factor,从而获取到不同大小维度的图像token。

整体来看,Pixel-Shuffle技术的实现过程并不复杂,且具有较高的灵活性和可扩展性。它可以轻松地与其他图像处理技术相结合,形成更为强大和全面的图像处理流水线。此外,Pixel-Shuffle还具有良好的泛化能力,能够适用于不同类型的图像数据和深度学习模型,为图像token压缩领域带来了更为广阔的应用前景。

12.1.2  Cross-layer Token Fusion压缩

Cross-layer Token Fusion策略是一种高效的方法,它通过细致评估各个token对模型效率和准确性的综合贡献,巧妙地在特定网络层上实施token fusion,从而有效突破了传统模型的局限性。在多模态模型中,token的合并以及相似token的精准识别,均建立在余弦相似度的坚实基础之上,确保了融合的准确性和有效性。

Cross-layer Token Fusion的具体实现如下所示。

def bipartite_soft_matching(
        metric: torch.Tensor, r: int,
        class_token: bool = False, distill_token: bool = False,
    ):

    protected = 0
    if class_token:
        protected += 1
    if distill_token:
        protected += 1

    # 我们最多只能减少50%的令牌
    t = metric.shape[1]
    r = min(r, (t - protected) // 2)

    if r <= 0:
        return metric, metric

    with torch.no_grad():
        # 归一化相似度矩阵
        metric = metric / metric.norm(dim=-1, keepdim=True)
        a, b = metric[..., ::2, :], metric[..., 1::2, :]
        scores = a @ b.transpose(-1, -2)

        if class_token:
            scores[..., 0, :] = -math.inf
        if distill_token:
            scores[..., :, 0] = -math.inf

        # 找到匹配分数最高的令牌对
        node_max, node_idx = scores.max(dim=-1)
        edge_idx = node_max.argsort(dim=-1, descending=True)[..., None]

        unm_idx = edge_idx[..., r:, :]  # 未合并的令牌
        src_idx = edge_idx[..., :r, :]  # 要合并的令牌
        dst_idx = node_idx[..., None].gather(dim=-2, index=src_idx)

        if class_token:
            # 确保类别令牌在开始位置
            unm_idx = unm_idx.sort(dim=1)[0]

    def merge(x: torch.Tensor, mode="mean") -> torch.Tensor:
        """
        合并操作,将指定的令牌合并。
        """
        src, dst = x[..., ::2, :], x[..., 1::2, :]
        n, t1, c = src.shape
        unm = src.gather(dim=-2, index=unm_idx.expand(n, t1 - r, c))
        src = src.gather(dim=-2, index=src_idx.expand(n, r, c))
        dst = dst.scatter_reduce(-2, dst_idx.expand(n, r, c), src, reduce=mode)

        if distill_token:
            return torch.cat([unm[:, :1], dst[:, :1], unm[:, 1:], dst[:, 1:]], dim=1)
        else:
            return torch.cat([unm, dst], dim=1)

    def unmerge(x: torch.Tensor) -> torch.Tensor:
        """
        取消合并操作,恢复原始令牌。
        """
        unm_len = unm_idx.shape[1]
        unm, dst = x[..., :unm_len, :], x[..., unm_len:, :]
        n, _, c = unm.shape

        src = dst.gather(dim=-2, index=dst_idx.expand(n, r, c))

        out = torch.zeros(n, metric.shape[1], c, device=x.device, dtype=x.dtype)

        out[..., 1::2, :] = dst
        out.scatter_(dim=-2, index=(2 * unm_idx).expand(n, unm_len, c), src=unm)
        out.scatter_(dim=-2, index=(2 * src_idx).expand(n, r, c), src=src)

        return out

    return merge, unmerge

其中,bipartite_soft_matching函数的输入参数解释如下:[yx1] [晓王2] 

  1. metric:输入的hidden_state张量,尺寸为[batch, tokens, channels]。
  2. r:要移除的令牌数量,最多为总令牌数的50%。
  3. class_token(可选):布尔值,指示是否有类别令牌。默认为False。
  4. distill_token(可选):布尔值,指示是否有蒸馏令牌。默认为False。

在具体使用上,这个函数是一种用于缩减图像处理中的image_token令牌数量的方案。通过应用平衡匹配集(ToMe),它能够在减少令牌(token)数量的同时保持信息的完整性。下面是使用bipartite_soft_matching进行token缩减的代码示例:

# 创建一个随机embedding张量,用于模拟token被torch.Embedding类计算的结果
embedding = torch.randn(size=(2, 48, 312))

# 调用bipartite_soft_matching函数,指定要移除的令牌数量r为18,不包含类别令牌
merge, unmerge = bipartite_soft_matching(embedding, 18, class_token=False)

# 使用merge函数处理embedding
hidden_states = merge(embedding)

# 打印处理后的hidden_states的形状
print(hidden_states.shape)  # 结果:torch.Size([2, 30, 312]),其中30 = 48 - 18

12.1.3  AvgPool的token压缩

AvgPoolProjector是一种用于取代corss_attention的图像token压缩算法,AvgPoolProjector通过自适应平均池化技术,在保留关键视觉信息的同时,有效缩减了图片token的数量。这一方法不仅简化了模型的复杂度,还提升了训练效率,使得在有限的计算资源下也能实现高效的视觉-文本模态对齐。更重要的是,由于其无参特性,AvgPoolProjector避免了繁琐的参数调优过程,且在实际应用中表现出色,为视觉与语言的跨模态理解提供了强有力的工具。

class AvgPoolProjector(nn.Module):

    def __init__(
            self,
            layer_num: int = 2,
            query_num: int = 36,    #这里是输出的seq_length
            mm_hidden_size: int = 384,  #图片经过patch_embedding后的d_model,也就是输入的维度
            llm_hidden_size: int = 384, #语言模型的d_model,也就是输出的维度
    ):
        super().__init__()
        self.layer_num = layer_num
        self.query_num = query_num
        self.mm_hidden_size = mm_hidden_size
        self.llm_hidden_size = llm_hidden_size
        self.build_net()

    def build_net(self):
        hw = int(self.query_num ** 0.5)
        sampler = nn.AdaptiveAvgPool2d((hw, hw))
        self.sampler = sampler
        modules = [nn.Linear(self.mm_hidden_size, self.llm_hidden_size)]
        for _ in range(1, self.layer_num):
            modules.append(nn.GELU())
            modules.append(nn.Linear(self.llm_hidden_size, self.llm_hidden_size))
        modules.append(torch.nn.RMSNorm(self.llm_hidden_size))
        self.mlp_projector = nn.Sequential(*modules)

    def forward(self, visual_feat: torch.Tensor):
        batch_size, seq_len, h_dim = visual_feat.shape  # 576
        # 计算平方根并断言它必须是整数
        root = seq_len ** 0.5
        assert float(int(root)) == root, f"{seq_len}的平方根不是整数"

        hw = int(seq_len ** 0.5)  # 24
        shaped_visual_feat = rearrange(visual_feat, "b (h w) d -> b d h w", h=hw,w=hw)  # torch.Size([64, 1024, 24, 24])
        pooled_visual_feat = self.sampler(shaped_visual_feat)  # torch.Size([64, 1024, 12, 12])
        reshaped_visual_feat = rearrange(pooled_visual_feat, "b d h w -> b (h w) d")  # [64, 144, 1024]
        output_feat = self.mlp_projector(reshaped_visual_feat)  # [64, 144, 4096])
        return output_feat

其具体使用示例如下:

embedding = torch.randn(size=(2, 49, 384))

avg_pool = AvgPoolProjector()

pooled = avg_pool(embedding)

print(pooled.shape)

AvgPoolProjector的好处显而易见:它能够在减少计算负担的同时,保持模型的性能。通过直接在patch级别进行下采样,它有效地避免了语义信息的损失,确保了视觉特征与文本之间的准确对应。此外,其实现的简洁性也大大提升了模型的易用性和可扩展性。

总的来说,AvgPoolProjector以简洁高效的方式实现了视觉token的压缩,不仅提升了模型的训练效率和性能,还为视觉与文本的跨模态交互提供了新的可能性。无论是在资源受限的场景下,还是在追求高性能的应用中,AvgPoolProjector都展现出了其独特的优势和潜力。

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值