DeepSeek大模型高性能核心技术与多模态融合开发 - 商品搜索 - 京东

除了我们前面自定义的基于注意力的视频分类模型，torchvision也自带了视频分类模型，并提供了模型的预训练参数。本节将基于这个预训练的视频分类模型来完成HMDB的动作分类。

14.3.1  torchvision简介

torchvision是PyTorch的一个图形图像库，专门服务于PyTorch深度学习框架，用于构建计算机视觉模型。它提供了丰富的功能和工具，帮助开发人员和研究人员轻松处理图像数据，从而加速计算机视觉应用的开发和部署。

在torchvision库中，有几个核心组件值得一提。首先是torchvision.datasets，这个模块包含了许多加载数据的函数以及常用的数据集接口，如MNIST、CIFAR10、ImageNet等，使得数据准备变得简单快捷。通过这些接口，用户可以轻松地下载、加载和预处理这些数据集，为模型训练做好准备。

另一个重要组件是torchvision.models，它提供了大量预训练的模型结构，如AlexNet、VGG、ResNet等。这些模型已经在大型数据集上做过训练，并可以直接用于各种计算机视觉任务，如图像分类、目标检测等。此外，用户还可以根据自己的需求对这些预训练模型进行微调，以适应特定的应用场景。

torchvision.transforms也是一个不可或缺的模块，它提供了丰富的图像变换操作，如裁剪、旋转、归一化等。这些变换可以帮助用户增强数据集，提高模型的泛化能力。同时，torchvision.transforms还提供了Compose类，用于将多个变换操作串联起来，形成一条完整的图像处理流水线。

除了上述核心组件外，torchvision库还提供了其他有用的方法和工具，如torchvision.utils中的函数可以帮助用户更方便地处理图像数据。这些实用工具使得torchvision库成为一个功能全面、易于使用的计算机视觉库。

下面代码是一个简单的函数，目的是帮助我们实现对视频数据的读取与转化。

import PIL
import torch
import torchvision
import torchvision.transforms as transformers

def preprocess_video(video: str,n_frames: int = 16):

    # Reading the video file
    vframes, _, _ = torchvision.io.read_video(filename=video, pts_unit='sec', output_format='TCHW')

    vframes = vframes.type(torch.float32)
    vframes_count = len(vframes)

    skip_frames = max(int(vframes_count/16), 1)

    selected_frame = vframes[0].unsqueeze(0)

    for i in range(1, n_frames):
        selected_frame = torch.concat((selected_frame, vframes[i * skip_frames].unsqueeze(0)))
    selected_resized_frame = trans(selected_frame)

    return selected_resized_frame

这段代码模仿了视频中随机抽取特定帧窗口的方法，首先获取视频总的帧数，然后根据定义的n_frames数值，在视频中截取相应的帧窗口数，作为视频数据集使用。[yx1] [晓王2] 

库orchvision库还提供了预训练模型供我们读取视频时使用，mvit_v2_s就是一个专用于视频分类的模型，其提供了预训练参数。整体mvit_v2_s的结构如图14-2所示。

图14-2  分类MViTv2-S整体结构

从上图可以看到。首先视频经Patch Partition（cube1）模块对输入视频进行分块和reshape；然后拼接CLS；后续的scale2、scale3、scale4和scale5使用Multi Head Pooling Attention（MHPA）逐步下采样时空分辨率的同时增加通道维度，每个阶段由n个Transformer块（MultiscaleBlock）组成，且只在scale3、scale4、scale5阶段的第一块中下采样时空分辨率，同时增加通道维度，scale2中MHPA的头数h=1（嵌入维度d较小），后面阶段的h均是前一阶段的2倍。

下面代码提供了一个使用torchvision导入模型和预训练参数的方法。需要注意，torchvision中预训练参数还提供了适配模型的维度变换工具，即对输入参数进行转换的函数。

weights = MViT_V2_S_Weights.DEFAULT

transforms = weights.transforms()

model = mvit_v2_s(weights=weights)

上面代码中，transforms = weights.transforms()是对维度进行变换的方法，简单说就是将原始的维度整合成一个新的维度，并用于模型计算。此外还需要注意，对于第一次使用预训练模型的读者来说，需要下载对应的模型参数，如图14-3所示。

图14-3  torchvision的数据准备

下面代码通过传入的transforms模块，在输出数据的同时，对数据的结构进行相应的变换处理。

def set_seed(seed: int = 929):
    np.random.seed(seed)
    torch.manual_seed(seed)
    random.seed(seed)

class HumanActionDataset(Dataset):

    def __init__(self, avi_files_list, label_list, n_frame = 16, transform = None):

        self.avi_files_list = avi_files_list
        self.label_list = label_list
        self.n_frame = n_frame
        self.transform = transform
        self.resize_trans = torchvision.transforms.Resize((224,224))

    def __len__(self):
        return len(self.label_list)

    def __getitem__(self, index):
        video_path = self.avi_files_list[index]

        # Reading the video file
        vframes, _, _ = torchvision.io.read_video(filename=video_path, pts_unit='sec', output_format='TCHW')
        vframes = vframes.type(torch.float32)
        vframes_count = len(vframes)

        # Selecting frames at certain interval
        skip_frames = max(int(vframes_count / self.n_frame), 1)
        selected_frame = vframes[0].unsqueeze(0)

        # Creating a new sequence of frames upto the defined sequence length.
        for i in range(1, self.n_frame):
            selected_frame = torch.concat((selected_frame, vframes[i * skip_frames].unsqueeze(0)))

        # Video label as per the classes list.
        label = torch.tensor(self.label_list[index])

        selected_frame = self.resize_trans(selected_frame)
        # Applying transformation to the frames
        if self.transform:
            return self.transform(selected_frame), label
        else:
            return selected_frame, label

上面代码在划分了训练集与测试集的基础上，主要完成了模型数据的准备和预处理，包括使用torchvision.io.read_video读取数据，并完成从中进行随机截取的任务。而transform用于使数据能够被调整适配mvit_v2_s模型的格式，并将其输出。

14.3.3  基于torchvision的端到端视频分类实战

我们可以直接使用torchvision提供的预训练模型来完成端到端的视频分类，完整代码如下所示。

import torch,torchvision
from torchvision.models.video import mvit_v2_s, MViT_V2_S_Weights
import einops

def create_model(num_classes: int, device: torch.device):

    weights = MViT_V2_S_Weights.DEFAULT
    transforms = weights.transforms()
    model = mvit_v2_s(weights=weights)
  
    dropout_layer = model.head[0]
    in_features = model.head[1].in_features
    model.head = torch.nn.Sequential(
        dropout_layer,
        torch.nn.Linear(in_features=in_features, out_features=num_classes, bias=True, device=device))
    return model.to(device), transforms

在上面代码中，我们首先搭建了一个完整的端到端模型框架，然后对输出端进行相应的替换操作，以确保其能够与我们预先定义的数据类别数目完美对接。此外，我们还采用了预定义的transforms维度处理模块，该模块在模型构建过程中被整合进来，并随模型一同返回，以便于后续的图像处理和模型应用。

完整的训练代码如下所示。

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

import pretrain_model
device = "cuda"

model,transforms = pretrain_model.create_model(num_classes=16,device=device)

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

BATCH_SIZE = 9
import get_data

if __name__ == '__main__':

    train_dataset = get_data.HumanActionDataset(get_data.avi_files_train, get_data.label_train,transform=transforms)
    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE,shuffle=True,num_workers=3)

    test_dataset = get_data.HumanActionDataset(get_data.avi_files_test, get_data.label_test,transform=transforms)
    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}%")


        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}%")

读者可以自行运行代码验证结果。