优化指南

Tiny-DL-Inference 的性能优化技术。

---

概述

Tiny-DL-Inference 实现了多种优化策略，以在 WebGPU 上最大化推理性能：

1. 算子融合 - 减少内存流量
2. 内存布局优化 - NCHW vs NHWC
3. 零拷贝操作 - 高效的张量视图
4. Im2Col 算法 - 卷积优化
5. 工作组调优 - GPU 占用优化

---

算子融合

概念

算子融合将多个操作合并为单个 GPU 内核，消除中间内存读/写。

示例：卷积 + 偏置 + ReLU

无融合：

内存：读取 → 卷积 → 写入 (1)
        读取 → 偏置 → 写入 (2)
        读取 → ReLU → 写入 (3)
        
总计：6 次内存操作

融合：

内存：读取 → 卷积+偏置+ReLU → 写入

总计：2 次内存操作

结果：3 倍内存带宽减少

何时使用融合算子

场景	建议
卷积 → 偏置 → ReLU	✅ 使用 Conv2dBiasReLUOperator
生产推理	✅ 始终使用融合
内存受限	✅ 关键优化
调试	❌ 使用分离算子
自定义激活	❌ 如可能手动融合

实现

// ❌ 低效：3 个独立内核
const conv = new Conv2dOperator(context);
const bias = ...; // 手动偏置添加
const relu = new ReLUOperator(context);

// ✅ 高效：1 个融合内核
const fused = new Conv2dBiasReLUOperator(context);
const output = await fused.forward([input, weight, bias], params);

---

内存布局

NCHW vs NHWC

图像数据的两种常见内存布局：

布局	格式	框架	缓存适配
NCHW	[N, C, H, W]	PyTorch	通道操作
NHWC	[N, H, W, C]	TensorFlow	空间操作

当前实现

卷积和池化：仅 NCHW

此选择优化：

• 卷积中的顺序通道访问
• 滤波器操作的缓存局部性

布局转换

使用不同格式时：

// 从 NHWC 转换为 NCHW 进行处理
const nchwTensor = await nhwcTensor.convertLayout('NCHW');

// 处理...
const output = await conv2d.forward([nchwTensor, weight], params);

// 如需要，转换回原始格式
const result = await output.convertLayout('NHWC');

性能影响

操作	开销
布局转换	高（完整张量复制）
NCHW 卷积	最佳
NHWC 卷积	不支持（先转换）

建议：尽可能在整个管道中保持数据为 NCHW。

---

零拷贝操作

张量视图

reshape() 方法创建共享底层 GPU 缓冲区的视图：

const tensor = new Tensor(context, [1, 3, 224, 224]);

// 零拷贝重塑
const flat = tensor.reshape([1, 150528]);
// flat.buffer === tensor.buffer（相同的 GPU 缓冲区）

优点

指标	复制	视图
时间	O(N)	O(1)
内存	2×	1×
GPU 开销	高	无

用例

1. Flatten 操作

   // Flatten 层（零开销）
   const flat = convOutput.reshape([batch, -1]);

2. 形状适配

   // 调整张量以适应不同的层期望
   const adapted = tensor.reshape([newBatch, newChannels, newH, newW]);

3. 批次操作

   // 重塑批次维度
   const batched = tensor.reshape([batchSize, -1, channels]);

限制

• 视图张量不能调整大小超出原始缓冲区
• 视图与父级共享生命周期（销毁父级 = 丢失数据）
• 布局转换需要实际数据移动

---

Im2Col 算法

概念

Im2Col 将卷积转换为矩阵乘法（GEMM），可以高度优化。

转换

输入:  [N, C, H, W]
         ↓ Im2Col
矩阵: [N*outH*outW, C*kH*kW]

权重: [K, C, kH, kW]
         ↓ 重塑
矩阵: [K, C*kH*kW]

输出 = GEMM(权重矩阵, Im2Col(输入))
        [K, N*outH*outW]
         ↓ 重塑
        [N, K, outH, outW]

实现工具

import { im2col } from 'tiny-dl-inference';

// 将图像转换为列格式
const col = im2col(input, {
  kernelHeight: 3,
  kernelWidth: 3,
  stride: 1,
  padding: 1
});

何时使用

场景	建议
大卷积核（5×5, 7×7）	考虑 Im2Col
自定义 GEMM 实现	需要 Im2Col
标准 3×3 卷积	直接卷积（当前）
分组卷积	专用算法

权衡

因素	直接卷积	Im2Col
内存	较低	较高（列缓冲区）
小卷积核	更快	开销
大卷积核	较慢	更快
实现	复杂	更简单

---

延迟资源销毁

问题

过早的缓冲区销毁会导致崩溃：

// ❌ 危险：在 GPU 使用缓冲区之前销毁
commandEncoder.copyBufferToBuffer(buffer, ...);
device.queue.submit([commandEncoder.finish()]);
buffer.destroy(); // 可能崩溃！

解决方案

deferDestroy() 在 GPU 工作完成后调度销毁：

// ✅ 安全：延迟销毁
commandEncoder.copyBufferToBuffer(buffer, ...);
device.queue.submit([commandEncoder.finish()]);
context.deferDestroy(buffer); // 安全！

实现

class GPUContext {
  private pendingCleanup: Set<Promise<void>> = new Set();
  
  deferDestroy(buffer: GPUBuffer | null): void {
    if (!buffer) return;
    
    const cleanup = this.waitForSubmittedWork()
      .then(() => buffer.destroy())
      .catch(() => { /* 忽略 */ });
    
    this.pendingCleanup.add(cleanup);
    cleanup.finally(() => this.pendingCleanup.delete(cleanup));
  }
  
  async sync(): Promise<void> {
    await this.waitForSubmittedWork();
    await Promise.allSettled([...this.pendingCleanup]);
    this.pendingCleanup.clear();
  }
}

---

基准测试

内置基准工具

import { Benchmark } from 'tiny-dl-inference';

const benchmark = new Benchmark();

const result = await benchmark.measureOperator(
  operator,      // 算子实例
  inputs,        // 输入张量
  params,        // 算子参数
  100            // 迭代次数
);

console.log({
  mean: result.meanMs,      // 平均执行时间
  stdDev: result.stdDevMs,  // 标准差
  min: result.minMs,        // 最小时间
  max: result.maxMs         // 最大时间
});

解读结果

指标	健康范围	高时的操作
平均值	一致	基线
标准差	< 平均值的 10%	检查变化
最小值	接近平均值	良好
最大值	< 2 倍平均值	检查异常值

分析技巧

1. 预热：测量前运行几次迭代

   for (let i = 0; i < 5; i++) {
     await operator.forward(inputs, params);
   }
   // 现在测量

2. 隔离变量：一次测试一个算子

3. 多次运行：跨多个基准会话取平均

4. 变化输入大小：使用真实数据大小测试

---

性能检查清单

部署前

• [ ] 适用时使用融合算子
• [ ] 最小化布局转换
• [ ] 使用真实输入分析
• [ ] 在目标硬件上测试
• [ ] 验证内存清理
• [ ] 与基线对比基准测试

运行时优化

• [ ] 尽可能重用张量
• [ ] 批处理多个推理
• [ ] 对临时变量使用 deferDestroy()
• [ ] 避免不必要的下载
• [ ] 预分配缓冲区

代码模式

// ✅ 良好：预分配输出
const output = new Tensor(context, outputShape);
for (const input of inputs) {
  await operator.forward([input, output], params);
}

// ❌ 糟糕：循环中分配
for (const input of inputs) {
  const output = new Tensor(context, outputShape); // 慢！
  await operator.forward([input], params);
}

---

硬件特定调优

工作组大小

默认工作组大小为 256。最优大小因 GPU 而异：

GPU 架构	最佳工作组大小
NVIDIA (Compute 7.0+)	256 或 512
AMD RDNA	128 或 256
Intel Xe	256
Apple M1/M2	256 或 512

测试工作组性能

// 基准测试不同的工作组大小
for (const wgSize of [128, 256, 512]) {
  // 修改着色器以使用 wgSize
  const result = await benchmarkWithWorkgroupSize(wgSize);
  console.log(`工作组 ${wgSize}: ${result.meanMs}ms`);
}

---

优化总结

技术	影响	工作量
算子融合	高	低（使用融合算子）
零拷贝视图	高	低（重塑 vs 复制）
布局一致性	中	中
延迟清理	高	低（自动）
缓冲区重用	中	中
工作组调优	中高	高

优先级：

1. 使用算子融合
2. 利用零拷贝视图
3. 保持一致的布局
4. 分析和调整工作组