注意力视频分类实战

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

上一节完成了视频数据集的准备，为接下来的实战打下了坚实的基础。本节中，我们将进一步探索，设计一种基于注意力架构的视频分类实战方案，并借助上一节自定义的数据准备形式，对视频进行精准分类。

在具体实现上，对于注意力模型而言，关键的一步在于如何将原始视频数据转换成一种模型能够高效处理的嵌入表示。这种嵌入表示不仅需要捕捉视频中的时序信息，还要能够突出关键帧和特征，以供注意力机制进行选择和聚焦。

为了达到这一目的，我们采用先进的深度学习技术，结合视频数据的特性来构建专门的嵌入层。这一层负责将视频帧序列转换为高维的特征向量，同时保留视频中的动态信息和空间结构。通过这些特征向量，注意力模型将能够更准确地识别视频中的关键内容，从而提升分类的准确度和效率。

在接下来的实战中，我们将详细阐述如何构建这种嵌入表示，并将其与注意力模型紧密结合，共同完成视频分类任务。

14.2.1  对于视频的Embedding编码器

对于视频的Embedding编码器设计，可以借鉴在2D图像处理中广泛应用的patch_embedding编码器思路。通过类似的方式，我们将视频数据划分为一系列时空块（spatio-temporal patches），每个块都包含了视频中的局部时空信息。

具体来说，我们首先将视频帧进行切片，生成一系列包含连续帧的小块。这样做不仅保留了视频中的时间连续性，还使得模型能够更有效地捕捉视频中的动态变化。接下来，将这些时空块通过Embedding层进行转换，生成对应的特征向量。这些特征向量将作为注意力模型的输入，用于后续的分类任务。

通过这种方式，我们能够充分利用视频数据的时空特性，同时降低模型的计算复杂度。此外，通过调整时空块的大小和数量，我们还可以进一步平衡模型的表达能力和计算效率，以适应不同场景下的视频分类任务。

下面代码是作者完成的一个视频Embedding编码器。

import torch
from einops.layers.torch import Rearrange,Reduce

def pair(t):
    return t if isinstance(t, tuple) else (t, t)

class ViT3D(torch.nn.Module):
    def __init__(self, image_size, image_patch_size, frames, frame_patch_size, dim, pool='cls', channels=3):
        super().__init__()  # 调用父类的初始化方法

        # 将输入的图像大小转换为高度和宽度
        image_height, image_width = pair(image_size)  # 例如: (128, 128)

        # 将输入的图像块大小转换为高度和宽度
        patch_height, patch_width = pair(image_patch_size)  # 例如: (16, 16)

        # 断言确保图像的高度和宽度可以被块的高度和宽度整除
        assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'

        # 断言确保帧数可以被帧块大小整除
        assert frames % frame_patch_size == 0, 'Frames must be divisible by frame patch size'  # 例如: 16帧, 每2帧一个块

        # 计算总的块数（考虑图像和帧的维度）
        num_patches = (image_height // patch_height) * (image_width // patch_width) * (frames // frame_patch_size)

        # 计算每个块的维度（考虑通道数、块的高度、宽度和帧块大小）
        patch_dim = channels * patch_height * patch_width * frame_patch_size  # 例如: 3*16*16*2=1536

        # 断言确保池化类型是'cls'（类标记）或'mean'（平均池化）
        assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'

        # 定义从输入到块嵌入的序列模型
        self.to_patch_embedding = torch.nn.Sequential(
            # 重新排列张量的维度以适应3D视频数据
            Rearrange('b (f pf) (h p1) (w p2) c -> b (f h w) (p1 p2 pf c)', p1=patch_height, p2=patch_width, pf=frame_patch_size),
            # 对重新排列后的数据进行层归一化
            torch.nn.RMSNorm(patch_dim),
            # 线性变换，将块维度转换为指定的隐藏维度（例如: 1536 -> 1024）
            torch.nn.Linear(patch_dim, dim),
        )

    # 定义前向传播方法
    def forward(self, x):
        # 将输入数据通过定义的序列模型，得到块嵌入
        x = self.to_patch_embedding(x.float())
        return x  # 返回处理后的数据

在上面代码中，我们分别将帧与维度进行拆分和重新组合，并对修正了隐藏维度获得了返回值。

14.2.2  视频分类模型的设计

接下来，我们需要完成视频分类模型的设计工作。在此过程中，我们采用经典的多头注意力模型MHA作为特征编码的核心组件，以协助我们完成视频的分类任务。多头注意力模型因其强大的特征提取和表示能力，在诸多序列处理任务中表现出色，我们相信它同样能在视频分类领域发挥重要作用。代码如下所示。

import torch
from torch import nn
import einops

from rotary_embedding_torch import RotaryEmbedding
class MultiHeadAttention(torch.nn.Module):
    def __init__(self, d_model, attention_head_num):
        super(MultiHeadAttention, self).__init__()
        self.attention_head_num = attention_head_num
        self.d_model = d_model

        assert d_model % attention_head_num == 0
        self.scale = d_model ** -0.5
        self.softcap_value = 50.
        self.per_head_dmodel = d_model // attention_head_num

        self.qkv_layer = torch.nn.Linear(d_model, 3 * d_model)
        self.out_layer = torch.nn.Linear(d_model, d_model)

        self.rotary_emb = RotaryEmbedding(dim = self.per_head_dmodel)

        "----------------------------------"
        self.q_scale = torch.nn.Parameter(torch.ones(self.per_head_dmodel))
        self.k_scale = torch.nn.Parameter(torch.ones(self.per_head_dmodel))

    def forward(self, embedding,past_length = 1024):
        b,l,d = embedding.shape

        qky_x = self.qkv_layer(embedding)
        q, k, v = torch.split(qky_x, split_size_or_sections=self.d_model, dim=-1)
        q = einops.rearrange(q, "b s (h d) -> b h s d", h=self.attention_head_num)
        k = einops.rearrange(k, "b s (h d) -> b h s d", h=self.attention_head_num)
        v = einops.rearrange(v, "b s (h d) -> b h s d", h=self.attention_head_num)

        q = torch.nn.functional.normalize(q, dim = -1) * self.q_scale * self.scale
        k = torch.nn.functional.normalize(k, dim = -1) * self.k_scale * self.scale

        q = self.rotary_emb.rotate_queries_or_keys(q)
        k = self.rotary_emb.rotate_queries_or_keys(k)

        sim = einops.einsum(q, k, 'b h i d, b h j d -> b h i j')

        i, j = sim.shape[-2:]

        attn = sim.softmax(dim=-1)
        out = einops.einsum(attn, v, 'b h i j, b h j d -> b h i d')
        embedding = einops.rearrange(out, "b h s d -> b s (h d)")
        embedding = self.out_layer(embedding)
        embedding = embedding[:,-l:]
        return embedding


from st_moe_pytorch import MoE,SparseMoEBlock
class ResidualAttention(nn.Module):
    """
    Residual Attention Block
    """

    def __init__(self,d_model,attention_head_num):

        super().__init__()
        self.attention_head_num = attention_head_num
        self.merge_norm = torch.nn.RMSNorm(d_model)
        self.attn = MultiHeadAttention(d_model, attention_head_num)  # self attention
        self.mlp = torch.nn.Sequential(torch.nn.GLU(),torch.nn.Linear((d_model // 2), d_model, bias=False))  # 注意这里输入的维度不要乘2

    def forward(self,x: torch.Tensor):
        residual = x

        x = self.merge_norm(x)
        attn_output = self.attn(x)
        x = residual + self.mlp(attn_output)  # norm and apply residual FFN
        return x

在上面代码中，我们采用了一个标准的注意力模型，作为视频分类任务的注意力基础计算框架，由于是对视频进行Embedding计算，我们去掉了因果掩码。

下面代码就是在注意力基础上完成的视频分类模型。

import blocks
class ViderClassificationModel_V1(torch.nn.Module):
    def __init__(self,dim = 384,head_num = 6,device = "cuda"):
        super(ViderClassificationModel_V1, self).__init__()
        self.patch_embedding_3d = ViT3D(112,16,frames=48,frame_patch_size=16,dim=dim)

        self.layers = []
        for _ in range(4):
            block = blocks.ResidualAttention(dim, head_num).to(device)
            self.layers.append(block)
        self.conv_layers = torch.nn.Sequential(
            torch.nn.Conv1d(147, 64, kernel_size=3, padding=1),
            torch.nn.RMSNorm(dim),
            torch.nn.Linear(dim,dim//2),
            torch.nn.Conv1d(64,32,kernel_size=3,padding=1),
        )
        self.logits_layer = torch.nn.Linear(6144,16)

        self.position_embedding = torch.nn.Parameter(torch.randn(size=(147,dim)),requires_grad=True)

    def forward(self, x):
        x = self.patch_embedding_3d(x) + self.position_embedding

        for block in self.layers:
            x = block(x)
        # torch.Size([6, 294, 384])

        x = self.conv_layers(x)
        x = torch.nn.Flatten()(x)
        x = torch.nn.functional.dropout(x,p = 0.1)
        x = self.logits_layer(x)

        return x

从代码中可以看到，这里就是一个比较简单的分类模型，首先通过patch_embedding对视频进行重新编码，之后注意力模型对特征进行计算，并最终通过logits_layer对结果进行分类计算。

14.2.3  视频分类模型的训练与验证

对于视频分类模型的训练与验证，我们可以借鉴经典的分类模型做法，采用交叉熵来计算损失函数，并最终返回相应的结果。代码如下所示。

import math
from tqdm import tqdm
import torch
from torch.utils.data import DataLoader

import model
device = "cuda"
model = model.ViderClassificationModel_V1(device=device)

model.to(device)
save_path = "./saver/video_classic.pth"
#model.load_state_dict(torch.load(save_path),strict=False)

BATCH_SIZE = 12

import get_data
train_dataset = get_data.TrainDataset()
train_loader = (DataLoader(train_dataset, batch_size=BATCH_SIZE,shuffle=True,num_workers=6))

test_dataset = get_data.TestDataset()
test_loader = (DataLoader(test_dataset, batch_size=BATCH_SIZE,shuffle=True))

optimizer = torch.optim.AdamW(model.parameters(), lr = 2e-5)
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,T_max = 1200,eta_min=2e-7,last_epoch=-1)
criterion = torch.nn.CrossEntropyLoss()

for epoch in range(128):
    model.train()
    pbar = tqdm(train_loader,total=len(train_loader))
    for frames_stack,label in pbar:
        frames_stack = frames_stack.to(device)
        label = label.to(device)
        logits = model(frames_stack)
        loss = criterion(logits,label)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        lr_scheduler.step()  # 执行优化器

        _, predicted = torch.max(logits.detach(), 1)  # 获取预测结果
        total = label.size(0)  # 获取当前批次的总样本数
        correct = (predicted == label).sum().item()  # 累加正确预测的样本数
        accuracy = 100 * correct / total  # 计算正确率

        pbar.set_description(f"epoch:{epoch + 1}, train_loss:{loss.item():.5f}, lr:{lr_scheduler.get_last_lr()[0] * 1000:.5f}, accuracy:{accuracy:.2f}%")

    if (epoch + 1)%3 == 0:
        torch.save(model.state_dict(), save_path)
        # 训练循环结束后的测试代码
        model.eval()  # 将模型设置为评估模式
        total_test = 0  # 测试集总样本数
        correct_test = 0  # 测试集正确预测样本数

        with torch.no_grad():  # 不需要计算梯度，节省内存和计算资源
            pbar_test = tqdm(test_loader, total=len(test_loader))
            for frames_stack, label in pbar_test:
                frames_stack = frames_stack.to(device)
                label = label.to(device)
                logits = model(frames_stack)
                _, predicted = torch.max(logits, 1)  # 获取预测结果
                total_test += label.size(0)  # 累加测试集的总样本数
                # 累加测试集正确预测的样本数
                correct_test += (predicted == label).sum().item()  
        accuracy_test = 100 * correct_test / total_test  # 计算测试集的正确率
        pbar_test.set_description(f"Test Accuracy: {accuracy_test:.2f}%")

        # 输出最终测试准确率
        print(f"Final Test Accuracy: {accuracy_test:.2f}%")

请读者自行训练与测试。