视频理解模型的推理优化：从帧采样策略到时空注意力机制的 GPU Kernel 定制

一、视频理解不是"逐帧图像识别"

视频理解的直观方案是：每秒抽取 1 帧，送入图像分类模型，对帧序列的结果取平均。这样做的问题很快暴露——视频中的核心信息往往不在一帧之内。一个人"拿起来杯子"这个动作，光看单帧只能看到手在杯子附近，看不到运动方向。时空信息丢失了。

视频 Transformer（如 TimeSformer、VideoMAE）通过时空注意力（Spatial-Temporal Attention）同时建模帧内和帧间的关系。但计算开销随之暴涨：一张 224×224 的视频帧约产生 196 个 patch（16×16 patch size），16 帧视频就是 3136 个 token。对于标准 Self-Attention（O(n²) 复杂度），3100 token 的注意力矩阵需要约 40MB 显存——仅一层。

帧采样策略直接影响延迟和准确度的折中。均匀采样（每 N 帧取一帧）简单但可能错过关键帧（如快速移动时的运动模糊帧恰是关键信息）。自适应采样（根据光流大小动态调整采样密度）更准确但引入了额外的光流计算开销。

另一个实际工程问题是视频预处理流水线。FFmpeg 解码一帧 H.265 视频约 2-4ms，但 PyTorch DataLoader 的 GIL 使得多线程解码的加速比受限于 2-3 倍。CPU 解码 + GPU 推理的流水线如果设计不当，GPU 的空闲等待时间占总体延迟的 40%。

二、帧采样策略与时空注意力的优化方案

帧采样策略的三种方案：

均匀采样是最简单的方案——16 帧从 10 秒视频中均匀抽取，每 0.625 秒一帧。CPU 开销低，但可能错过 0.1 秒内的快速动作。

自适应采样要求先对视频做一次快速光流分析（使用稀疏光流，如 Lucas-Kanade），标记运动幅度超过阈值的帧为关键帧。在这些关键帧前后增加采样密度。这种方法增加了 10-20% 的前处理延迟，但识别准确度提升 5-8%。

关键帧提取依赖场景切换检测——两帧之间的颜色直方图差异超过阈值时标记为场景切换。对于监控视频（大部分时间画面静止），这种方法可将处理帧数减少 80%。

时空注意力的 Divided Attention 优化：

标准 Joint Attention（所有 patch 互相看）的复杂度是 O((T×S)²)，其中 T 是时间帧数，S 是空间 patch 数。Divided Attention 先将注意力计算拆分为空间维度和时间维度：

• Spatial Attention: 每帧内部，patch 之间互相注意。S 个 patch → O(S²) 每帧。
• Temporal Attention: 同一空间位置的 patch，跨帧互相注意。T 帧 → O(T²) 每位置。

总复杂度从 O((T×S)²) 降至 O(T×S² + S×T²)。当 T=16, S=196 时，计算量从 (16×196)² ≈ 9.8M 降至 16×196² + 196×16² ≈ 0.66M——减少约 93%。

三、视频 Transformer 的推理优化实现

use std::sync::Arc;
use rayon::prelude::*;
use candle_core::{Device, Tensor, DType};

/// 视频帧 —— 解码后的单帧
#[derive(Clone)]
pub struct VideoFrame {
    /// 帧时间偏移（秒）
    pub timestamp: f32,
    /// RGB 像素数据 [3 × H × W]，已做 normalize
    pub pixels: Tensor,
}

/// 帧采样器 —— 决定从视频中选取哪些帧
pub struct FrameSampler {
    /// 目标采样帧数
    pub num_frames: usize,
    /// 输入分辨率
    pub input_size: (usize, usize), // (H, W)
    /// patch 大小
    pub patch_size: usize,
}

impl FrameSampler {
    /// 均匀采样：从帧序列中均匀选取 num_frames 帧
    pub fn uniform_sample(&self, frames: &[VideoFrame]) -> Vec<VideoFrame> {
        if frames.len() <= self.num_frames {
            // 帧数不够时，重复最后一帧填充（实际应做 padding）
            let mut result = frames.to_vec();
            while result.len() < self.num_frames {
                result.push(frames.last().unwrap().clone());
            }
            return result;
        }
        
        let step = frames.len() as f64 / self.num_frames as f64;
        let mut sampled = Vec::with_capacity(self.num_frames);
        
        for i in 0..self.num_frames {
            let idx = (i as f64 * step) as usize;
            // clamp: 防止浮点舍入误差导致索引越界
            let idx = idx.min(frames.len() - 1);
            sampled.push(frames[idx].clone());
        }
        
        sampled
    }

    /// 自适应采样：基于运动检测的关键帧采样
    /// 
    /// 算法：计算相邻帧的像素差，累计运动量。
    /// 当累计运动量超过阈值时，选取当前帧。
    /// 策略：运动剧烈区域密集采样，静止区域稀疏采样。
    pub fn adaptive_sample(
        &self,
        frames: &[VideoFrame],
        motion_threshold: f64,
    ) -> Vec<VideoFrame> {
        let mut sampled = Vec::new();
        let mut motion_accum = 0.0;
        
        // 第一帧始终采样
        if let Some(first) = frames.first() {
            sampled.push(first.clone());
        }
        
        for i in 1..frames.len() {
            // 计算像素差 —— 简化：用 Tensor 的 MSE
            // 实际应使用光流（Optical Flow）来精确衡量运动
            let diff = frames[i].pixels.sub(&frames[i-1].pixels)
                .unwrap_or_else(|_| Tensor::zeros((1,), DType::F32, &Device::Cpu).unwrap());
            let motion = diff.abs()?.mean_all()?.to_scalar::<f64>()?;
            
            motion_accum += motion;
            
            // 累计运动量超过阈值或在均匀采样位置时采样
            if motion_accum >= motion_threshold 
                || sampled.len() < self.num_frames 
                    && i % (frames.len() / self.num_frames) == 0 
            {
                sampled.push(frames[i].clone());
                motion_accum = 0.0; // 重置
            }
            
            if sampled.len() >= self.num_frames {
                break;
            }
        }
        
        // 确保帧数达标
        if sampled.len() < self.num_frames {
            // 从均匀采样位置补齐
            self.uniform_sample(frames)
        } else {
            sampled
        }
    }
}

/// 视频 Patch 嵌入 —— 将帧分割为 patch 并嵌入
pub struct VideoPatchEmbedding {
    pub patch_size: usize,
    pub embed_dim: usize,
}

impl VideoPatchEmbedding {
    /// 将视频帧序列转换为 patch tokens
    /// 
    /// 输入: [T, 3, H, W] → 输出: [T, N, D]
    /// T = 帧数, N = (H/p) × (W/p) = patch 数, D = embed_dim
    pub fn embed(
        &self,
        frames: &[VideoFrame],
        device: &Device,
    ) -> Result<Tensor, candle_core::Error> {
        let t = frames.len();
        
        // 为所有帧生成 patch tokens
        let tokens: Vec<Tensor> = frames.iter().map(|frame| {
            // 简化：每帧 [3, H, W] → [N, D]
            // 实际实现：
            // 1. 用 Conv2d(patch_size, patch_size, stride=patch_size) 完成 patch 提取
            // 2. 展平 patch → 线性投影 → 加位置编码
            Tensor::zeros(
                (1, (self.input_size.0 / self.patch_size) * (self.input_size.1 / self.patch_size), self.embed_dim),
                DType::F32,
                device,
            )
        }).collect::<Result<Vec<_>, _>>()?;
        
        // 堆叠为 [T, N, D]
        Tensor::stack(&tokens.iter().map(|t| t).collect::<Vec<_>>(), 0)
    }
}

/// Divided Attention: 空间注意力 + 时间注意力
pub struct DividedAttention {
    pub embed_dim: usize,
    pub num_heads: usize,
    /// QKV 投影权重 —— 空间注意力和时间注意力共享或独立
    pub spatial_qkv: Arc<Tensor>,
    pub temporal_qkv: Arc<Tensor>,
}

impl DividedAttention {
    /// 前向传播：先做空间注意力，再做时间注意力
    /// 
    /// 输入: [B, T, N, D] (Batch × Frames × Patches × Embed)
    /// 输出: [B, T, N, D]
    pub fn forward(
        &self,
        x: &Tensor,
    ) -> Result<Tensor, candle_core::Error> {
        let (b, t, n, d) = (x.dim(0)?, x.dim(1)?, x.dim(2)?, x.dim(3)?);
        
        // ===== 1. 空间注意力 (Spatial Attention) =====
        // 将 [B, T, N, D] reshape 为 [B×T, N, D]
        // 每帧内部的 patch 之间进行注意力计算
        let x_spatial = x.reshape((b * t, n, d))?;
        
        // QKV 投影
        let q = x_spatial.matmul(&self.spatial_qkv)?; // 简化：应分开 QKV
        let k = x_spatial.matmul(&self.spatial_qkv)?;
        let v = x_spatial.matmul(&self.spatial_qkv)?;
        
        // Scaled Dot-Product Attention
        let scale = (d as f64).sqrt() as f32;
        // QK^T / sqrt(d)
        let attn_weights = q.matmul(&k.transpose(1, 2)?)? / scale;
        // Softmax + Dropout（简化实现）
        let attn_weights = candle_nn::ops::softmax(&attn_weights, 2)?;
        // Weighted Sum
        let spatial_out = attn_weights.matmul(&v)?;
        
        // 恢复形状: [B×T, N, D] → [B, T, N, D]
        let spatial_out = spatial_out.reshape((b, t, n, d))?;
        
        // ===== 2. 时间注意力 (Temporal Attention) =====
        // 将 [B, T, N, D] reshape 为 [B×N, T, D]
        // 同一空间位置的 patch 跨帧进行注意力计算
        let x_temporal = spatial_out.permute((0, 2, 1, 3))?
            .reshape((b * n, t, d))?;
        
        let q_t = x_temporal.matmul(&self.temporal_qkv)?;
        let k_t = x_temporal.matmul(&self.temporal_qkv)?;
        let v_t = x_temporal.matmul(&self.temporal_qkv)?;
        
        let attn_weights_t = q_t.matmul(&k_t.transpose(1, 2)?)? / scale;
        let attn_weights_t = candle_nn::ops::softmax(&attn_weights_t, 2)?;
        let temporal_out = attn_weights_t.matmul(&v_t)?;
        
        // 恢复形状: [B×N, T, D] → [B, N, T, D] → [B, T, N, D]
        let temporal_out = temporal_out.reshape((b, n, t, d))?
            .permute((0, 2, 1, 3))?;
        
        Ok(temporal_out)
    }
}

/// 视频理解推理管线
pub struct VideoInferencePipeline {
    pub sampler: FrameSampler,
    pub embedder: VideoPatchEmbedding,
    pub attention: DividedAttention,
    pub device: Device,
}

impl VideoInferencePipeline {
    /// 端到端推理
    pub fn infer(&self, video_path: &str) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
        // 1. 视频解码与帧采样
        // 实际实现需要接入 FFmpeg 进行解码
        let raw_frames: Vec<VideoFrame> = self.decode_video(video_path)?;
        let sampled = self.sampler.uniform_sample(&raw_frames);
        
        // 2. Patch 嵌入
        let tokens = self.embedder.embed(&sampled, &self.device)?;
        
        // 3. Divided Attention 编码
        let encoded = self.attention.forward(&tokens)?;
        
        // 4. 全局平均池化 → 分类头
        // 实际实现：mean over T and N dimensions, then linear classifier
        let pooled = encoded.mean_all()?;
        let result = pooled.to_vec1::<f32>()?;
        
        Ok(result)
    }

    fn decode_video(&self, _path: &str) -> Result<Vec<VideoFrame>, Box<dyn std::error::Error>> {
        // 使用 ffmpeg-next 解码视频流
        // 返回所有帧（或按帧率选取的关键帧）
        Ok(vec![])
    }
}

// 并行多视频处理：利用 Rayon 数据并行
pub fn batch_process_videos(pipeline: &VideoInferencePipeline, video_paths: &[String]) 
    -> Vec<Result<Vec<f32>, Box<dyn std::error::Error>>> 
{
    video_paths.par_iter()
        .map(|path| pipeline.infer(path))
        .collect()
}

关键设计决策：

• Divided Attention 的 reshape 操作：[B, T, N, D] → [B×T, N, D] 将多帧的空间 patch 并排处理，允许批量矩阵乘法。在 GPU 上，大尺寸矩阵乘法比多个小尺寸矩阵乘法更高效。
• 自适应采样中的 motion_accum：累计运动量而非单帧运动量作为采样触发条件——避免因场景瞬变导致的过度采样。这本质上是一个低通滤波器。
• Arc<Tensor> 用于 QKV 权重：多个视频可以共享同一组权重，不必为每次推理拷贝权重矩阵。

四、视频推理优化的适用边界与权衡

适用场景：

• 固定时长视频分析（10-30 秒），如短视频理解、监控事件检测。
• 模型架构使用 Transformer 且帧数 ≥ 8 的场景——Divided Attention 的优势随帧数增加而更明显。
• GPU 推理——reshape 操作在 GPU 上几乎是免费的（只是改变 stride），适合批量处理。

不适用场景：

• 实时视频流处理（延迟 < 100ms）。Divided Attention 需要完整的多帧才能计算，引入了帧缓冲延迟。
• 帧数 < 4 的短片段。此时 Divided Attention 的优势不明显，Joint Attention 更简单。
• CPU-only 推理——视频 Transformer 在 CPU 上的推理时间可能达到秒级。

主要权衡：

1. Divided Attention 的精度损失：空间-时间分解假设空间和时间信息是可分离的。对于精细的时空交互（如物体遮挡），分解带来约 1-2% 的准确度损失。
2. 自适应采样的计算开销：光流计算本身是密集运算。如果采样节省的时间小于光流计算增加的时间，自适应采样反而更慢。
3. Patch 大小的选择：16×16 的 patch 是大模型的标准配置。对于低分辨率视频（如 112×112），patch 过大会导致 patch 数量太少，注意力无法有效捕捉细节。

五、总结

1. Divided Attention 将时空联合注意力的 O((T×S)²) 降至 O(T×S² + S×T²)，在 16 帧 196 patch 场景下计算量减少约 93%。
2. 均匀采样是帧采样策略的性能基线——当光流计算开销超过节省的推理时间时，自适应采样反而不划算。
3. reshape 操作实现空间-时间分离注意力的批量计算，是 GPU 高效推理的关键技术。
4. 运动累计（Motion Accumulation）而非瞬时运动量作为自适应采样的触发条件，避免了高频噪声导致的过度采样。
5. 视频 Transformer 推理优化的三个杠杆：帧采样率（减少帧数）、注意力分解（减少计算量）、patch 大小（减少 token 数）。