低光图像增强：Transformer轻量化与RK3588部署优化

1. 项目背景与核心挑战

低光环境下的图像增强一直是计算机视觉领域的硬骨头。传统方法往往面临噪声放大、色彩失真和细节丢失三大难题。去年我们在RK3399上尝试部署基于CNN的增强模型时，就遇到了实时性不足和内存溢出的问题——处理一张1080P图像需要近2秒，这在实际安防监控场景中根本无法接受。

Transformer架构的出现带来了转机。其自注意力机制能更好地建模全局依赖关系，特别适合处理低光图像中常见的非均匀光照问题。但标准Transformer的计算复杂度与图像尺寸呈平方关系，直接部署到RK3588这类边缘设备上依然不现实。

2. 模型轻量化关键技术

2.1 混合注意力机制设计

我们采用了一种混合窗口注意力方案：

python复制class MixedAttention(nn.Module):
    def __init__(self, dim, window_size=8):
        super().__init__()
        self.window_attn = WindowAttention(dim, window_size)
        self.global_attn = SimplifiedGlobalAttention(dim)
        
    def forward(self, x):
        # 局部窗口注意力处理细节
        local_feat = self.window_attn(x)
        # 全局注意力处理光照关系
        global_feat = self.global_attn(x)
        return local_feat + global_feat

这种设计在PSNR指标上仅比纯全局注意力低0.3dB，但内存占用减少了62%。实测在4K图像上，推理速度从原来的17fps提升到43fps。

2.2 动态通道裁剪策略

通过分析不同光照条件下各通道的激活强度，我们开发了动态通道裁剪模块：

1. 在模型训练时加入通道重要性评估层
2. 部署时根据设备剩余内存动态调整通道保留比例
3. 极端情况下可自动切换为8bit量化模式

这个方案使得模型在RK3588上内存占用稳定在1.2GB以内，避免了因内存波动导致的进程崩溃。

3. RK3588部署优化实践

3.1 NPU加速适配要点

RK3588的NPU对Transformer支持有限，需要特别注意：

• 将LayerNorm替换为GroupNorm
• 避免使用动态形状的reshape操作
• 最大程度利用NPU的INT8加速能力

我们开发的自动优化工具能完成以下转换：

bash复制python convert.py --model enhancer.onnx \
                 --output enhancer_rk3588.rknn \
                 --quantize --npu-optimize

3.2 内存管理技巧

通过实测发现三个关键经验：

1. 使用rknn-toolkit的memory_profile模式分析峰值内存
2. 设置GPU/NPU共享内存池（至少预留200MB）
3. 启用CMA连续内存分配器减少碎片

4. 实测性能对比

在自建的低光测试集（含2000张各类场景图像）上：

5. 典型问题排查指南

问题1：NPU推理结果异常

• 检查onnx模型中是否含有不支持的操作（如动态切片）
• 验证各层输入/输出尺度是否超出NPU限制

问题2：内存泄漏

• 使用valgrind工具检查内存分配
• 特别注意OpenCV后处理中的临时buffer

问题3：图像边缘伪影

• 调整padding策略为'reflection'
• 在模型末端添加轻量级去伪影模块

这个项目最深的体会是：边缘设备部署永远要在模型精度和工程可行性之间找平衡点。我们最终方案虽然PSNR比学术SOTA低了1.2dB，但实现了27fps的实时处理能力，这才是工业落地的关键。