EDVR是视频超分辨率领域一个经典的模型，它凭借PCD（Pyramid, Cascading and Deformable convolution）对齐模块和TSA（Temporal and Spatial Attention）融合模块在多个基准测试上取得了领先效果。不过在实际部署或者训练时，很多开发者都会遇到显存不足的困扰，尤其是处理高分辨率视频序列时，GPU内存占用会迅速飙升。这篇文章会从EDVR的显存消耗源头出发，介绍一种名为“金字塔级联”的结构改造方法，并讨论如何通过轻量化特征对齐进一步压缩显存占用。

EDVR显存瓶颈到底在哪里

要优化显存，首先得搞清楚EDVR为什么吃显存。原始EDVR的整体流程是：输入连续多帧低分辨率图像，每帧先经过一个共享的特征提取网络得到基础特征，然后相邻帧的特征会送入PCD对齐模块，对齐到参考帧。对齐后的特征经过TSA模块进行时空注意力融合，最后送入重建网络输出高分辨率图像。显存占用最大的部分通常集中在对齐模块和融合模块。

PCD对齐模块内部包含金字塔特征提取和级联可变形卷积。金字塔部分会在多个尺度上计算特征，每个尺度都要保留中间结果用于后续级联，这相当于同时存储了多份不同分辨率的特征图。级联可变形卷积又会生成额外的偏移量和调制掩码，这两者的通道数往往与特征图相同甚至更大，直接导致显存翻倍。以输入五帧、每帧特征通道64、空间分辨率128×128为例，PCD模块在一个尺度上就可能产生数百MB的中间张量。

TSA融合模块则需要对对齐后的多个特征计算时间注意力和空间注意力。时间注意力通过卷积压缩通道后计算相似度矩阵，空间注意力则使用金字塔池化获取全局上下文。这些操作本身参数不多，但会产生大量的临时激活值，尤其是在融合五帧特征时，显存开销随帧数线性增长。另外，重建网络如果采用多个RDB（Residual Dense Block），其密集连接也会存储许多中间特征。

一个容易忽视的显存来源是自动梯度。训练时PyTorch默认会保存所有中间激活值用于反向传播，即使代码里没有显式存储，这些张量也会占用显存。EDVR的对齐和融合模块中间激活值特别多，所以训练时的显存占用往往是推理时的三到五倍。

金字塔级联如何降低显存

金字塔级联的核心思路很简单：不要在最高分辨率上一次性完成对齐和融合，而是先在低分辨率尺度上做粗对齐和融合，然后逐步上采样到更高分辨率做细化。这样做最直接的好处就是中间特征的空间尺寸大幅减小，显存占用与分辨率面积成正比，所以降到原来一半的分辨率就能节省约四分之三的显存。

具体实现上，可以在输入帧经过浅层特征提取后，通过下采样得到多个尺度的特征，比如1/4、1/2和原始分辨率。先在最粗的1/4尺度上运行PCD对齐和TSA融合，得到粗融合特征。然后把这个粗融合特征上采样到1/2尺度，与该尺度下的原始特征拼接或者相加，再送入一个轻量级的对齐和融合模块进行细化。最后在原始分辨率上再做一次细化。这种做法相当于把原本单个大模块的计算拆分成多级小模块，每一级只处理对应尺度的特征，显存峰值可以控制在较小范围内。

代码示例展示一个两级的金字塔级联结构，用PyTorch实现：

import torch
import torch.nn as nn
import torch.nn.functional as F
class PyramidCascadeEDVR(nn.Module):
    def __init__(self, nf=64, nframes=5):
        super().__init__()
        self.nf = nf
        self.nframes = nframes
        # 浅层特征提取
        self.conv_first = nn.Conv2d(3, nf, 3, 1, 1)
        # 粗尺度对齐（1/4分辨率）
        self.coarse_align = PCDAlign(nf, groups=4)
        self.coarse_fusion = TSAFusion(nf, nframes)
        # 细尺度对齐（1/2分辨率）
        self.fine_align = PCDAlign(nf, groups=2)
        self.fine_fusion = TSAFusion(nf, nframes)
        # 上采样与重建
        self.upsample = nn.Sequential(
            nn.Conv2d(nf, nf * 4, 3, 1, 1),
            nn.PixelShuffle(2),
            nn.LeakyReLU(0.1, True),
            nn.Conv2d(nf, nf * 4, 3, 1, 1),
            nn.PixelShuffle(2),
            nn.LeakyReLU(0.1, True),
            nn.Conv2d(nf, 3, 3, 1, 1)
        )
    def forward(self, x):
        # x: [B, T, C, H, W]
        B, T, C, H, W = x.shape
        x = x.view(B * T, C, H, W)
        feat = self.conv_first(x)  # [B*T, nf, H, W]
        feat = feat.view(B, T, -1, H, W)
        # 生成1/4尺度特征
        feat_coarse = F.interpolate(feat.view(B*T, -1, H, W), scale_factor=0.25, mode='bilinear', align_corners=False)
        feat_coarse = feat_coarse.view(B, T, -1, H//4, W//4)
        # 粗对齐与融合
        aligned_coarse = self.coarse_align(feat_coarse)
        fused_coarse = self.coarse_fusion(aligned_coarse)  # [B, nf, H//4, W//4]
        # 上采样到1/2尺度
        fused_half = F.interpolate(fused_coarse, scale_factor=2, mode='bilinear', align_corners=False)
        # 原始1/2尺度特征
        feat_half = F.interpolate(feat.view(B*T, -1, H, W), scale_factor=0.5, mode='bilinear', align_corners=False)
        feat_half = feat_half.view(B, T, -1, H//2, W//2)
        ref_feat_half = feat_half[:, self.nframes//2]  # 参考帧特征
        combined_half = torch.cat([fused_half, ref_feat_half], dim=1)
        combined_half = nn.Conv2d(self.nf*2, self.nf, 1, 1, 0).to(x.device)(combined_half)
        # 细对齐与融合
        aligned_half = self.fine_align(torch.cat([feat_half[:, :self.nframes//2], combined_half.unsqueeze(1), feat_half[:, self.nframes//2+1:]], dim=1))
        fused_half = self.fine_fusion(aligned_half)
        # 上采样到原始分辨率
        fused_full = F.interpolate(fused_half, scale_factor=2, mode='bilinear', align_corners=False)
        out = self.upsample(fused_full)
        return out

上述代码只是一个简化示例，实际使用时需要替换PCDAlign和TSAFusion为对应的模块，并且注意各个模块的输入输出形状。金字塔级联的另一个好处是可以复用低分辨率计算得到的偏移量，在高分辨率细化时只需要预测残差偏移，这样也能减少参数和计算量。

需要指出的是，金字塔级联并不是简单地降低分辨率就万事大吉。如果所有尺度都使用同样复杂的PCD和TSA模块，整体参数量不会减少，只是显存峰值降低，但总计算量可能因为多级处理而增加。因此在细化阶段应该采用更轻量的模块，比如减少可变形卷积的层数、减小通道数，或者用普通的卷积替代部分可变形卷积。

特征对齐的轻量化改造

原始EDVR的PCD对齐模块包含一个多尺度特征金字塔，每个尺度都使用可变形卷积进行对齐，然后级联起来。这种设计虽然精度高，但可变形卷积在计算偏移量时需要额外的卷积层，而且可变形卷积本身在PyTorch中实现时会产生大量中间张量。轻量化的思路是保留对齐的核心能力，但减少冗余计算和显存。

一种可行的方案是用分组可变形卷积替代普通可变形卷积。分组可变形卷积把通道分成若干组，每组独立计算偏移量，可以大幅减少偏移量预测的参数量和计算量。例如原来使用64通道的普通可变形卷积，偏移量预测卷积输出128通道（x和y方向各64），如果改为4组，每组16通道，那么偏移量预测输出只有32通道，显存直接降为原来的四分之一。代码中只需要在DeformConv2d的参数中设置groups即可，不过需要注意自定义可变形卷积实现是否支持分组。

另一种方案是使用光流估计网络替代可变形卷积做对齐。光流估计网络通常比较轻量，比如使用简单的U-Net结构预测运动矢量，然后通过warp操作对齐特征。warp操作不需要存储额外的偏移量和掩码，显存占用更低。虽然光流对齐在遮挡和大运动场景下精度可能略低于可变形卷积，但配合金字塔级联可以在粗尺度用可变形卷积、细尺度用光流，取长补短。

以下是一个轻量级光流对齐模块的示例代码：

import torch
import torch.nn as nn
import torch.nn.functional as F
class LightFlowAlign(nn.Module):
    def __init__(self, nf=64):
        super().__init__()
        self.conv1 = nn.Conv2d(nf*2, nf, 3, 1, 1)
        self.conv2 = nn.Conv2d(nf, nf, 3, 1, 1)
        self.flow_conv = nn.Conv2d(nf, 2, 3, 1, 1)
    def forward(self, ref_feat, neigh_feat):
        # ref_feat: [B, nf, H, W], neigh_feat: [B, nf, H, W]
        x = torch.cat([ref_feat, neigh_feat], dim=1)
        x = F.relu(self.conv1(x), inplace=True)
        x = F.relu(self.conv2(x), inplace=True)
        flow = self.flow_conv(x)  # [B, 2, H, W]
        # warp操作
        aligned_feat = self.warp(neigh_feat, flow)
        return aligned_feat
    def warp(self, x, flow):
        B, C, H, W = x.size()
        # 生成采样网格
        grid_y, grid_x = torch.meshgrid(torch.arange(H, device=x.device), torch.arange(W, device=x.device))
        grid = torch.stack((grid_x, grid_y), dim=0).float()
        grid = grid.unsqueeze(0).repeat(B, 1, 1, 1)
        vgrid = grid + flow
        vgrid[:, 0, :, :] = 2.0 * vgrid[:, 0, :, :] / max(W - 1, 1) - 1.0
        vgrid[:, 1, :, :] = 2.0 * vgrid[:, 1, :, :] / max(H - 1, 1) - 1.0
        vgrid = vgrid.permute(0, 2, 3, 1)
        output = F.grid_sample(x, vgrid, align_corners=True)
        return output

这个光流对齐模块参数量少，显存占用低，特别适合金字塔级联的细尺度阶段。另外还可以在TSA融合模块中做文章，比如使用更小的空间注意力金字塔，或者把时间注意力从计算完整相似度矩阵改为分组注意力，都能进一步压缩显存。

特征对齐的轻量化不仅仅是换模块，还需要考虑训练策略。轻量化模块可能无法从零开始学到和原始PCD一样的对齐精度，因此可以在训练时使用知识蒸馏，让原始EDVR作为教师网络，轻量级网络学习对齐后的特征分布。或者在训练初期使用较大的学习率，配合余弦退火，帮助轻量模块快速收敛。

工程实践中的显存优化组合策略

除了模型结构层面的金字塔级联和特征对齐轻量化，还可以结合一些工程技巧进一步降低显存。梯度检查点（Gradient Checkpointing）是一种常用的方法，它在前向传播时只保存部分中间激活值，反向传播需要时重新计算。对于EDVR这种显存大户，在PCD和TSA模块上应用梯度检查点可以节省大量显存，代价是增加约20%到30%的训练时间。PyTorch中可以使用torch.utils.checkpoint轻松实现。

混合精度训练也是显存优化的利器。使用float16存储参数和激活值，显存占用直接减半，同时利用Tensor Core还能加速计算。需要注意EDVR中可变形卷积和光流warp操作对数值精度比较敏感，建议在这些模块保留float32，只对特征提取和重建部分使用混合精度。

如果仍然遇到显存不足，可以考虑降低输入帧数。EDVR通常使用5帧，减少到3帧可以节省将近一半的时序相关显存。或者采用循环推理的方式，一次只处理两帧，逐步累计融合，但这样会损失一定的对齐精度。

在实际部署时，还可以把金字塔级联拆分成多阶段执行，比如先在GPU上完成粗尺度对齐融合，然后把中间结果转移到CPU内存，再在GPU上处理细尺度。虽然这会增加数据搬运开销，但对于显存极小的设备来说是一个可行的折中方案。

最后谈谈训练与推理的不同侧重。训练时需要保存激活值用于反向传播，所以显存优化重点在于梯度检查点、混合精度和小批量。推理时不需要保存激活值，显存压力主要来自模型参数和IFrame缓存，此时金字塔级联可以结合算子融合（比如将卷积和激活融合）来减少峰值显存。针对特定硬件，还可以使用TensorRT或ONNX Runtime进行图优化，消除冗余内存分配。

总的来说，解决EDVR显存不足需要从模型结构和工程实现两个层面同时入手。金字塔级联为结构优化提供了明确方向，特征对齐轻量化则在不明显损失精度的情况下压缩显存。结合梯度检查点、混合精度等措施，完全可以在8GB甚至更小显存的显卡上跑通EDVR的训练和推理。