BERT预训练模型实现图片分类：从原理到工程实践

1次阅读

没有评论

共计 2214 个字符，预计需要花费 6 分钟才能阅读完成。

在传统的图片分类任务中，CNN（卷积神经网络）一直是主流选择。但随着 Transformer 架构在 NLP 领域的成功，越来越多的研究者开始探索其在视觉领域的应用。那么，为什么我们需要 BERT 这样的语言模型来处理视觉数据呢？

CNN 的局限性：虽然 CNN 在局部特征提取上表现优异，但在处理长距离依赖关系时不如 Transformer 高效。
BERT 的优势：BERT 的注意力机制（Attention Mechanism）能够捕捉全局上下文信息，这在某些复杂的视觉任务中可能更有优势。
核心挑战：
维度匹配：BERT 的输入是词向量序列，而图片是二维矩阵，如何将图片转换为 BERT 可处理的序列？
位置编码：在 NLP 中，位置编码用于表示词的位置信息，但在视觉任务中，这种编码方式需要调整。
计算效率：图片的分辨率通常较高，直接输入 BERT 会导致计算量激增。

ViT 和 BEiT 是目前最流行的视觉 Transformer 方案，它们都通过 Patch Embedding 将图片分块并转换为序列。

ViT：直接将图片分块，每个块通过线性映射得到嵌入向量。
BEiT：在 ViT 的基础上引入了掩码图像建模（Masked Image Modeling），类似于 BERT 的掩码语言建模。

Patch Embedding 是 BERT 处理图片的关键步骤。具体来说：

将图片划分为固定大小的块（如 16×16 像素）。
每个块通过线性映射转换为一个向量（类似于词向量）。
将所有块的向量拼接成一个序列，作为 BERT 的输入。

在视觉任务中，位置编码需要重新设计。传统的正弦位置编码可能不适合图片的二维结构。一种常见的做法是使用可学习的位置编码，或者在 Patch Embedding 中加入二维位置信息。

以下是使用 PyTorch 实现的完整代码示例：

import torch
import torch.nn as nn
from transformers import BertModel, BertConfig

class BERTForImageClassification(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        config = BertConfig.from_pretrained('bert-base-uncased')
        self.bert = BertModel(config)
        self.patch_embedding = nn.Conv2d(3, config.hidden_size, kernel_size=16, stride=16)
        self.classifier = nn.Linear(config.hidden_size, num_classes)

    def forward(self, x):
        # 图像分块处理
        patches = self.patch_embedding(x)  # [batch_size, hidden_size, num_patches_h, num_patches_w]
        patches = patches.flatten(2).transpose(1, 2)  # [batch_size, num_patches, hidden_size]

        # BERT 输入
        outputs = self.bert(inputs_embeds=patches)
        pooled_output = outputs.pooler_output

        # 分类头
        logits = self.classifier(pooled_output)
        return logits

图像分块处理 ：nn.Conv2d 通过卷积操作将图片分块，并映射到 BERT 的隐藏层维度。flatten和 transpose 操作将分块后的图片转换为序列格式。
跨模态 Attention Mask：在 BERT 中，Attention Mask 用于控制哪些位置可以参与注意力计算。对于图片任务，通常不需要额外的 Mask。
微调学习率：建议使用较小的学习率（如 2e-5）进行微调，避免破坏预训练模型的权重。

输入分辨率对 GPU 显存的影响非常大。以下是不同分辨率下的显存占用曲线：