1. 为什么要在 Vulkan 上跑神经网络1.1 从一次部署翻车说起去年帮一个做工业质检的朋友处理现场问题他们的设备是一台国产化平台上的工控机显卡是某款国产 GPU驱动只提供了 Vulkan 接口CUDA 那套东西完全用不了。他们原本的 OCR 推理代码是基于 CUDA 写的迁移到这台机器上直接跑不起来。当时第一反应是换硬件但设备已经装到产线上了拆机成本太高。于是就有了这个方案用 Vulkan 把神经网络推理跑起来。这件事让我意识到一个被很多人忽略的事实GPU 计算不等于 CUDA。在桌面和数据中心场景CUDA 确实是事实标准但到了嵌入式、国产化、跨平台这些场景Vulkan 反而是更通用的选择。Vulkan 是 Khronos 组织维护的跨平台图形与计算 APIWindows、Linux、Android 都原生支持只要驱动实现了 Vulkan Compute理论上就能跑神经网络。这篇文章要讲的就是怎么用 C 和 Vulkan 从零搭一个神经网络推理引擎并且用它完整跑通一次 PP-OCR 的 GPU 推理。PP-OCR 是百度飞桨开源的 OCR 模型包含文本检测和文本识别两个阶段模型结构不算复杂但足够覆盖卷积、池化、激活、全连接这些常见算子拿来练手非常合适。1.2 谁适合看这篇内容如果你属于下面几类人这篇内容应该对你有用手上有非 NVIDIA 的 GPU想跑神经网络推理但被 CUDA 卡住的开发者做跨平台应用希望一套推理代码在 Windows、Linux、Android 上都能跑想理解 GPU 计算底层原理不满足于只会调 PyTorch 接口做国产化适配需要绕开 CUDA 生态的技术选型需要说明的是这篇不是 Vulkan 入门教程我不会从vkCreateInstance开始讲起。我默认你对 Vulkan 的基本概念Instance、Device、Queue、CommandBuffer、Pipeline有了解或者至少愿意边看边查。重点放在怎么把神经网络的计算映射到 Vulkan 的 Compute Pipeline 上以及实际跑 PP-OCR 时踩到的那些坑。1.3 整体技术路线先给一个全局视图避免后面迷路。整个方案分四层层级职责关键技术模型层解析 ONNX 模型提取计算图ONNX Runtime 的模型解析接口图优化层算子融合、内存复用规划自定义图遍历与调度执行层把算子映射到 Vulkan ComputeCompute Pipeline DescriptorSet数据层张量在 CPU/GPU 间搬运Staging Buffer Device Local Memory模型层用 ONNX 作为中间格式因为 PP-OCR 官方提供了 ONNX 导出省得自己从 Paddle 格式转。图优化层做的是把 ConvBNReLU 这类组合融合成一个算子减少 Kernel 启动次数。执行层是核心每个算子对应一个 SPIR-V Shader通过 Vulkan 的 Compute Pipeline 执行。数据层负责 Host 和 Device 之间的数据传输以及中间张量的显存分配。这个路线的好处是不依赖任何厂商特定的计算库只要 Vulkan 驱动正常代码就能跑。代价是要自己实现算子工作量不小但可控性极强。2. Vulkan Compute 与 CUDA 的本质差异2.1 编程模型的根本区别很多人第一次接触 Vulkan Compute 会觉得别扭因为它和 CUDA 的编程模型差别很大。CUDA 是我写一个 Kernel指定 grid 和 block剩下的交给驱动Vulkan 是我准备好所有资源录制命令提交到队列然后等结果。前者是命令式的后者是声明式的。具体来说CUDA 里你写__global__ void add(float* a, float* b, float* c, int n) { int i blockIdx.x * blockDim.x threadIdx.x; if (i n) c[i] a[i] b[i]; }然后addgrid, block(a, b, c, n)就完事了。Vulkan 里同样的加法你需要创建两个 Storage Buffer 存输入一个存输出写一个 GLSL Compute Shader编译成 SPIR-V创建 DescriptorSetLayout 和 PipelineLayout创建 Compute Pipeline分配 DescriptorSet绑定 Buffer录制 CommandBuffervkCmdDispatch提交到 Queue等待 Fence代码量差了一个数量级。但换来的是显式的资源管理和跨平台能力。CUDA 的隐式管理在 NVIDIA 卡上很爽换到别的卡就抓瞎。2.2 内存模型的差异CUDA 有统一内存Unified MemorycudaMallocManaged分配的指针 CPU 和 GPU 都能访问驱动自动做迁移。Vulkan 没有这个概念你必须显式管理Host Visible Host CoherentCPU 能直接读写但 GPU 访问慢适合小数据量频繁更新Device LocalGPU 访问最快但 CPU 不能直接访问需要通过 Staging Buffer 中转Host Visible Device Local某些集成显卡上存在CPU 和 GPU 都能访问但性能介于两者之间跑神经网络推理时权重和中间张量应该放在 Device Local 内存里输入输出数据通过 Staging Buffer 搬运。这个原则和 CUDA 里尽量减少 Host-Device 传输是一致的但 Vulkan 需要你手动做。我实测下来在独立显卡上Device Local 内存的带宽比 Host Visible 内存高 3 到 5 倍。PP-OCR 的检测模型大概 4MB 权重识别模型 10MB 左右如果权重放在 Host Visible 内存里每次推理都要走 PCIe 读一遍延迟会明显增加。2.3 同步机制的差异CUDA 里 Kernel 是异步启动的但同一个流里的 Kernel 按顺序执行你不太需要关心同步。Vulkan 里 CommandBuffer 提交后你需要用 Fence 或 Semaphore 来同步Fence用于 CPU 等待 GPU 完成Semaphore用于 GPU 内部不同 Queue 之间的同步Barrier用于同一个 CommandBuffer 内不同 Dispatch 之间的同步跑神经网络时如果两个算子之间有数据依赖必须在 CommandBuffer 里插入vkCmdPipelineBarrier确保前一个算子的写入对后一个算子的读取可见。这个 Barrier 如果漏了结果会随机出错而且很难调试因为有时候能跑对有时候跑错。提示Barrier 的 srcStageMask 和 dstStageMask 要设对。Compute Shader 之间的依赖用VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT访问类型用VK_ACCESS_SHADER_WRITE_BIT和VK_ACCESS_SHADER_READ_BIT。设错了要么性能差要么结果错。3. 把神经网络算子翻译成 Vulkan Shader3.1 算子拆解以 PP-OCR 的卷积层为例PP-OCR 的检测模型是 DBDifferentiable Binarization结构主干网络是 MobileNetV3 的变体里面大量使用深度可分离卷积。识别模型是 CRNN 结构包含卷积层和 LSTM 层。我们先从最基础的卷积说起。一个标准的 2D 卷积输入张量形状[N, C_in, H, W]权重[C_out, C_in, K, K]输出[N, C_out, H_out, W_out]。在 CUDA 里你可能会用 im2col GEMM 的方式或者直接写一个滑窗 Kernel。Vulkan 里我选择直接滑窗因为 im2col 会产生大量中间数据显存开销大。GLSL Compute Shader 的写法#version 450 layout(local_size_x 8, local_size_y 8, local_size_z 1) in; layout(binding 0) readonly buffer Input { float data[]; } input_tensor; layout(binding 1) readonly buffer Weight { float data[]; } weight; layout(binding 2) readonly buffer Bias { float data[]; } bias; layout(binding 3) writeonly buffer Output { float data[]; } output_tensor; layout(push_constant) uniform Params { int N, C_in, H, W; int C_out, K, stride, pad; int H_out, W_out; } p; void main() { int w_out int(gl_GlobalInvocationID.x); int h_out int(gl_GlobalInvocationID.y); int nc int(gl_GlobalInvocationID.z); int n nc / p.C_out; int c_out nc % p.C_out; if (w_out p.W_out || h_out p.H_out || n p.N) return; float sum bias.data[c_out]; for (int c_in 0; c_in p.C_in; c_in) { for (int kh 0; kh p.K; kh) { for (int kw 0; kw p.K; kw) { int h_in h_out * p.stride - p.pad kh; int w_in w_out * p.stride - p.pad kw; if (h_in 0 || h_in p.H || w_in 0 || w_in p.W) continue; int in_idx ((n * p.C_in c_in) * p.H h_in) * p.W w_in; int w_idx ((c_out * p.C_in c_in) * p.K kh) * p.K kw; sum input_tensor.data[in_idx] * weight.data[w_idx]; } } } int out_idx ((n * p.C_out c_out) * p.H_out h_out) * p.W_out w_out; output_tensor.data[out_idx] sum; }这个 Shader 每个线程算一个输出像素local size 设成 8x8。实测下来对于 3x3 卷积这个朴素实现已经能跑到可用水平但如果追求极致性能需要做共享内存优化把输入 tile 加载到 shared memory减少全局内存访问次数。3.2 深度可分离卷积的特殊处理MobileNetV3 里的深度可分离卷积分两步Depthwise Conv 和 Pointwise Conv。Depthwise 是每个通道独立卷积Pointwise 是 1x1 卷积。这两个算子的 Shader 和标准卷积略有不同。Depthwise 的权重形状是[C, 1, K, K]每个通道一个卷积核。Shader 里c_out和c_in是一一对应的不需要内层循环遍历所有输入通道。Pointwise 的 K1可以退化成矩阵乘法用 GEMM 的优化思路来做。我一开始把 Depthwise 也当成标准卷积来算结果性能差了一倍多。后来单独写了 Depthwise 的 Shader把C_in循环去掉速度立刻上来了。这个经验说明不要试图用一个通用 Kernel 覆盖所有情况针对特定算子做特化收益很明显。3.3 激活函数与逐元素算子ReLU、Sigmoid、HardSwish 这些激活函数都是逐元素的Shader 很简单#version 450 layout(local_size_x 256) in; layout(binding 0) readonly buffer Input { float data[]; } input_tensor; layout(binding 1) writeonly buffer Output { float data[]; } output_tensor; layout(push_constant) uniform Params { int total; } p; void main() { int i int(gl_GlobalInvocationID.x); if (i p.total) return; float x input_tensor.data[i]; // HardSwish: x * relu6(x 3) / 6 float relu6 clamp(x 3.0, 0.0, 6.0); output_tensor.data[i] x * relu6 / 6.0; }这类算子的瓶颈不在计算而在内存带宽。每个元素读一次写一次计算量可以忽略。所以优化方向是合并算子比如把 Conv Bias HardSwish 融合成一个 Shader中间结果不落显存。3.4 算子融合的收益与代价算子融合是推理优化的常规手段。在 Vulkan 里融合意味着把多个算子的逻辑写进同一个 Shader减少 Dispatch 次数和中间张量的显存读写。以 Conv BN ReLU 为例BN 在推理阶段可以折叠进卷积权重W W * gamma / sqrt(var eps) b (b - mean) * gamma / sqrt(var eps) beta折叠之后Conv 的权重和偏置直接更新BN 就消失了。ReLU 则可以在 Conv Shader 的最后一步加上sum max(sum, 0.0)。这样三个算子变成一个Dispatch 次数从 3 次降到 1 次中间张量也省了。但融合不是无脑做。我踩过一个坑把 Conv 和后面的 Add残差连接融合结果 Shader 里需要同时读两个输入张量寄存器压力变大occupancy 下降反而变慢了。后来改成不融合让两个算子各自 Dispatch性能更好。融合的收益要看具体算子和硬件不能一概而论。4. PP-OCR 推理的完整实现链路4.1 模型解析与图构建PP-OCR 的 ONNX 模型可以用 ONNX Runtime 的 C API 解析。核心接口是Ort::Session但我不需要它做推理只需要它帮我读出计算图Ort::Session session(env, model_path, session_options); Ort::Graph graph session.GetGraph(); // 伪代码实际用 ModelProto更直接的方式是用onnx.proto生成的 C 类直接读ModelProto遍历graph.node()。每个 Node 有op_type、input、output、attribute。我把这些信息转成自己的图结构struct TensorInfo { std::string name; std::vectorint64_t shape; size_t element_count; size_t byte_size; }; struct NodeInfo { std::string op_type; std::vectorstd::string inputs; std::vectorstd::string outputs; std::unordered_mapstd::string, Attribute attrs; }; struct Graph { std::vectorTensorInfo tensors; std::vectorNodeInfo nodes; std::unordered_mapstd::string, int tensor_index; };解析完之后做一次拓扑排序确保算子按依赖顺序执行。ONNX 的节点列表本身通常是拓扑有序的但为了保险还是排一次。4.2 显存分配与内存复用PP-OCR 检测模型有几十个中间张量如果每个都单独分配显存峰值占用会很大。实际上很多张量的生命周期不重叠可以复用同一块内存。我实现了一个简单的内存池先统计每个张量的首次使用和最后使用位置然后做区间调度把不重叠的张量分配到同一个 offset。具体做法是维护一个空闲块列表按大小排序分配时找最合适的块Best Fit。class MemoryPool { struct Block { size_t offset; size_t size; bool free; }; std::vectorBlock blocks; size_t total_size 0; public: size_t allocate(size_t size) { // Best fit 查找 for (auto b : blocks) { if (b.free b.size size) { b.free false; return b.offset; } } // 没有合适的扩展 size_t offset total_size; total_size size; blocks.push_back({offset, size, false}); return offset; } void free(size_t offset) { for (auto b : blocks) { if (b.offset offset) { b.free true; break; } } } };实测下来PP-OCR 检测模型的内存复用能把峰值显存从 200MB 降到 80MB 左右。这个优化在嵌入式设备上很关键因为显存往往只有 1GB 甚至更少。4.3 文本检测的后处理PP-OCR 检测模型的输出是一个概率图形状[1, 1, H, W]每个像素值表示该位置是文本的概率。后处理要做的是二值化概率大于阈值通常 0.3的像素置 1否则置 0连通域分析找出所有连通的文本区域轮廓提取对每个连通域提取外接矩形或多边形文本框过滤去掉太小或长宽比异常的框非极大值抑制合并重叠的框这些操作在 CPU 上做就行不需要 GPU。但要注意概率图从 GPU 读回 CPU 时用vkMapMemory映射 Staging Buffer拷贝完再vkUnmapMemory。如果频繁映射性能会受影响所以尽量一次性把整张图读回来。4.4 文本识别的序列处理识别模型是 CRNN输入是检测框裁剪出来的文本图像输出是字符序列。CRNN 包含 CNN 特征提取和 RNN 序列建模两部分。CNN 部分和检测模型类似RNN 部分在 ONNX 里通常导出为 LSTM 或 GRU。LSTM 在 Vulkan 里实现比较麻烦因为它是串行依赖的每个时间步依赖上一个时间步的隐藏状态。我的做法是把 LSTM 展开成多个 Dispatch每个时间步一个 Dispatch用 Barrier 同步。这样虽然 Dispatch 次数多但实现简单正确性有保证。如果追求性能可以用并行扫描Parallel Scan的算法把 LSTM 并行化但实现复杂度高很多。对于 PP-OCR 这种序列长度通常不超过 25 的场景串行展开已经够用了。4.5 端到端跑通的实测数据在一台搭载某款国产 GPU 的工控机上PP-OCR 的实测数据阶段耗时ms备注检测模型推理45输入 960x960检测后处理12CPU 上做识别模型推理8单条文本平均 20 字符识别后处理2CTC 解码总计67单张图对比同机器的 CPU 推理8 核 ARMGPU 版本快了大约 6 倍。虽然比不上 NVIDIA 显卡上 CUDA 的性能但在没有 CUDA 的环境里这个结果已经能用了。5. 踩过的坑与排查思路5.1 结果全零DescriptorSet 绑定错误第一次跑通 Pipeline 后输出全是零。排查过程检查 Shader 编译SPIR-V 生成正常检查 Buffer 创建大小和用法标志正确检查 DescriptorSet 更新vkUpdateDescriptorSets的dstBinding和 Shader 里的binding对不上问题出在第 3 步。我在 Shader 里写的是layout(binding 0)但 DescriptorSetLayout 里 binding 0 是 Inputbinding 1 是 Weightbinding 2 是 Output。更新 DescriptorSet 时我把 Output 绑到了 binding 1导致 Shader 读 Weight 的位置读到了输出 Buffer。提示DescriptorSet 的 binding 编号一定要和 Shader 里的layout(binding N)严格对应。建议用常量定义不要硬编码数字。5.2 结果随机错误Barrier 缺失修好绑定后结果大部分正确但偶尔有几个像素不对。这种随机错误最头疼因为不是每次都复现。用 RenderDoc 抓帧分析发现两个连续的 Dispatch 之间没有 Barrier。前一个 Dispatch 写输出 Buffer后一个 Dispatch 读同一个 Buffer但 GPU 的调度顺序不保证有时候后一个 Dispatch 先执行读到了旧数据。加上 Barrier 后问题消失VkMemoryBarrier barrier{}; barrier.sType VK_STRUCTURE_TYPE_MEMORY_BARRIER; barrier.srcAccessMask VK_ACCESS_SHADER_WRITE_BIT; barrier.dstAccessMask VK_ACCESS_SHADER_READ_BIT; vkCmdPipelineBarrier(cmd, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, 0, 1, barrier, 0, nullptr, 0, nullptr);这个坑的教训是只要两个 Dispatch 之间有数据依赖就必须加 Barrier。不要心存侥幸觉得应该没问题GPU 的调度顺序是不确定的。5.3 性能不达预期Occupancy 太低Barrier 加完后结果正确了但性能比预期慢很多。用厂商提供的 Profiler 一看Compute Shader 的 occupancy 只有 25%。原因是 local size 设得太大。我一开始把卷积 Shader 的 local size 设成 16x16想着一次多算点。但 16x16256 个线程每个线程用的寄存器又多导致一个 Compute Unit 上只能驻留很少的 Wavefrontoccupancy 上不去。改成 8x864 后occupancy 提升到 75%性能提升了近一倍。local size 不是越大越好要根据 Shader 的寄存器用量和硬件的 Compute Unit 规格来调。一般来说64 到 256 之间比较合适具体要实测。5.4 显存不足中间张量没复用在嵌入式设备上跑的时候显存只有 1GB检测模型跑到一半就 OOM 了。用vkAllocateMemory的返回值检查发现分配失败。原因是每个中间张量都单独分配了显存没有复用。加上前面说的内存池后峰值显存降下来了。另外权重数据在推理过程中是只读的可以放在一块独立的 Device Local 内存里和中间张量分开管理。还有一个细节Vulkan 的vkAllocateMemory有数量限制很多设备最多只允许 4096 次分配。如果每个张量都单独分配很容易超限。用内存池把多个张量合并到一次分配也能规避这个问题。6. 跨平台适配的注意事项6.1 不同厂商驱动的差异Vulkan 是标准但各家驱动实现有差异。我遇到过的情况某厂商驱动不支持VK_FORMAT_R32_SFLOAT的 Storage Image只能用 Storage Buffer某厂商驱动的vkCmdDispatch对 local size 有额外限制超过 256 会报错某厂商驱动的 Barrier 实现有 bug需要加额外的VK_ACCESS_MEMORY_READ_BIT应对策略是运行时查询设备能力根据VkPhysicalDeviceLimits和VkPhysicalDeviceFeatures调整实现。比如 Storage Image 不支持就全部用 Storage Bufferlocal size 超限就拆成多次 Dispatch。6.2 Android 上的特殊处理Android 的 Vulkan 环境和桌面不太一样内存分配更紧张需要更激进的内存复用Surface 和 Compute 共用 Queue要注意同步某些设备不支持VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT和HOST_VISIBLE的组合在 Android 上跑 PP-OCR我建议把模型量化成 FP16 甚至 INT8减少显存占用和带宽压力。Vulkan 支持 FP16 的 Shader需要VK_KHR_shader_float16_int8扩展但兼容性参差不齐要做好 fallback。6.3 精度问题的排查跨平台跑的时候发现同一个模型在不同设备上输出有细微差异。排查下来有几个原因浮点运算顺序不同GPU 的并行归约顺序和 CPU 不同导致浮点误差累积FP16 精度损失某些设备默认用 FP16 做中间计算精度不够Fast Math 优化驱动可能开启快速数学优化牺牲精度换速度对于 OCR 这种任务细微的数值差异通常不影响最终结果因为后处理有阈值和二值化。但如果做的是数值敏感的任务比如金融计算就要注意这些差异必要时强制用 FP32 并关闭 Fast Math。7. 这套方案还能怎么扩展7.1 支持更多算子目前实现了 Conv、Depthwise Conv、Pointwise Conv、Add、Mul、ReLU、HardSwish、Sigmoid、MaxPool、GlobalAvgPool、Reshape、Transpose、LSTM。要支持更多模型还需要补Concat多个输入张量拼接到一起注意对齐Split一个张量拆成多个和 Concat 相反Resize/Upsample双线性插值检测模型的上采样会用到Softmax识别模型的 CTC 解码前会用到Gemm全连接层可以用矩阵乘法的优化思路每加一个算子就是写一个 Shader 加一个 Dispatch 封装工作量可控。7.2 性能优化的方向当前实现是能跑离跑得快还有距离。后续可以做的优化Winograd 卷积3x3 卷积用 Winograd 算法可以减少乘法次数Shared Memory 优化把输入 tile 加载到 shared memory减少全局内存访问异步传输用多个 Queue 并行做数据传输和计算算子融合更激进的融合策略减少 Dispatch 和显存读写FP16/INT8 量化减少带宽压力提升计算吞吐这些优化每一个都能带来可观的性能提升但实现复杂度也高。建议先把功能跑通再逐步优化。7.3 和其他推理框架的对比市面上已经有一些支持 Vulkan 的推理框架比如 NCNN、MNN。它们的优势是算子库成熟、优化做得好缺点是定制性差遇到不支持的算子要改框架代码。自己实现的好处是完全可控想加什么算子就加什么算子想怎么优化就怎么优化。代价是要自己维护一套代码。选择哪条路取决于你的具体需求如果只是跑标准模型用现成框架更省事如果有特殊算子或特殊硬件自己实现更灵活。我在这个项目里的体会是Vulkan 跑神经网络不是能不能的问题而是值不值得的问题。如果你的目标平台只有 Vulkan那没得选只能自己实现。如果平台有 CUDA 或其他成熟的计算库优先用现成的把精力放在业务逻辑上。技术选型要看投入产出比不要为了技术而技术。