算子实现

Tiny-DL-Inference 中所有神经网络算子的完整指南。

算子概览

算子	描述	WGSL 实现
ReLU	线性整流单元	逐元素 `max(0, x)`
Softmax	归一化指数	数值稳定的 softmax
MaxPool	2D 最大池化	滑动窗口最大值
Conv2d	2D 卷积	直接卷积算法
Conv2dBiasReLU	融合卷积+偏置+ReLU	单内核优化
Flatten	张量重塑	零拷贝视图操作
Dense	全连接	矩阵乘法

ReLUOperator

线性整流单元激活函数。

描述

将所有负值设为零，保持正值不变。

公式

f(x) = max(0, x)

使用

typescript

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

const relu = new ReLUOperator(context);
const output = await relu.forward([input]);

WGSL 实现

wgsl

@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
  let idx = global_id.x;
  if (idx >= numElements) { return; }
  
  let x = input[idx];
  output[idx] = select(0.0, x, x > 0.0);
}

性能

内存：每个元素 1 次读取 + 1 次写入
计算：每个元素 1 次比较
适用：所有元素大小

SoftmaxOperator

将 logits 转换为概率分布。

描述

沿指定轴应用指数归一化。

公式

softmax(x_i) = exp(x_i - max(x)) / sum(exp(x_j - max(x)))

max(x) 减法确保数值稳定性。

参数

名称	类型	默认值	描述
`axis`	`number`	-1（最后）	计算 softmax 的轴

使用

typescript

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

const softmax = new SoftmaxOperator(context);

// 沿最后一维 softmax
const output = await softmax.forward([logits], { axis: -1 });

// 沿特定维度 softmax
const output = await softmax.forward([logits], { axis: 1 });

WGSL 实现

wgsl

// 两遍算法确保数值稳定性
// 第一遍：查找最大值
var maxVal = -3.402823e+38f;
for (var i = 0u; i < size; i = i + 1u) {
  maxVal = max(maxVal, input[i]);
}

// 第二遍：计算 exp(x - max) 和求和
var sum = 0.0;
for (var i = 0u; i < size; i = i + 1u) {
  let expVal = exp(input[i] - maxVal);
  output[i] = expVal;
  sum = sum + expVal;
}

// 第三遍：归一化
for (var i = 0u; i < size; i = i + 1u) {
  output[i] = output[i] / sum;
}

输出属性

所有值在范围 [0, 1] 内
值的总和等于 1.0
保持相对顺序

MaxPoolOperator

用于下采样的 2D 最大池化层。

描述

将输入划分为池化区域，输出每个区域的最大值。

参数

名称	类型	默认值	描述
`poolSize`	`number`	2	池化窗口大小（方阵）
`stride`	`number`	poolSize	池之间的步长

输出形状

对于输入 [N, C, H, W]：

output_height = floor((H - poolSize) / stride) + 1
output_width = floor((W - poolSize) / stride) + 1
output_shape = [N, C, output_height, output_width]

使用

typescript

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

const maxpool = new MaxPoolOperator(context);

// 2x2 池化，步长 2
const output = await maxpool.forward([input], { poolSize: 2, stride: 2 });

// 3x3 池化，步长 1
const output = await maxpool.forward([input], { poolSize: 3, stride: 1 });

WGSL 实现

wgsl

for (var ph = 0u; ph < outH; ph = ph + 1u) {
  for (var pw = 0u; pw < outW; pw = pw + 1u) {
    var maxVal = -3.402823e+38f;
    
    for (var kh = 0u; kh < poolSize; kh = kh + 1u) {
      for (var kw = 0u; kw < poolSize; kw = kw + 1u) {
        let h = ph * stride + kh;
        let w = pw * stride + kw;
        let idx = ((n * C + c) * H + h) * W + w;
        maxVal = max(maxVal, input[idx]);
      }
    }
    
    let outIdx = ((n * C + c) * outH + ph) * outW + pw;
    output[outIdx] = maxVal;
  }
}

性能

内存：顺序访问模式
计算：每次输出 O(poolSize²) 次比较
最佳：poolSize ∈

Conv2dOperator

用于特征提取的 2D 卷积层。

描述

在输入区域和内核之间应用滑动点积。

参数

名称	类型	默认值	描述
`channels`	`number`	必需	输出通道数（K）
`kernelSize`	`number`	3	卷积核大小（方阵）
`stride`	`number`	1	卷积之间的步长
`padding`	`number`	0	零填充大小

输入/输出形状

输入：[N, C, H, W]
权重：[K, C, kH, kW]
偏置（可选）：[K]

输出：[N, K, outH, outW]

其中：

outH = floor((H + 2*padding - kH) / stride) + 1
outW = floor((W + 2*padding - kW) / stride) + 1

使用

typescript

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

const conv2d = new Conv2dOperator(context);

const output = await conv2d.forward([input, weight, bias], {
  channels: 32,
  kernelSize: 3,
  stride: 1,
  padding: 1
});

WGSL 实现

wgsl

// 对于每个输出位置
for (var oh = 0u; oh < outH; oh = oh + 1u) {
  for (var ow = 0u; ow < outW; ow = ow + 1u) {
    var sum = 0.0;
    
    // 在所有输入通道和内核上卷积
    for (var ic = 0u; ic < C; ic = ic + 1u) {
      for (var kh = 0u; kh < kH; kh = kh + 1u) {
        for (var kw = 0u; kw < kW; kw = kw + 1u) {
          let ih = oh * stride + kh - padding;
          let iw = ow * stride + kw - padding;
          
          // 检查边界
          if (ih >= 0 && ih < H && iw >= 0 && iw < W) {
            let inputVal = input[((n * C + ic) * H + ih) * W + iw];
            let weightVal = weight[((oc * C + ic) * kH + kh) * kW + kw];
            sum = sum + inputVal * weightVal;
          }
        }
      }
    }
    
    // 添加偏置
    sum = sum + bias[oc];
    
    let outIdx = ((n * K + oc) * outH + oh) * outW + ow;
    output[outIdx] = sum;
  }
}

性能特征

卷积核大小	相对速度	用例
1×1	快	逐点变换
3×3	标准	标准卷积
5×5	较慢	更大的感受野
7×7	最慢	初始层

Conv2dBiasReLUOperator

融合的卷积+偏置+ReLU，实现最佳性能。

描述

将三个操作合并为单个内核，消除中间内存读/写。

公式

输出 = ReLU(Conv2d(输入, 权重) + 偏置)
     = max(0, Conv2d(输入, 权重) + 偏置)

性能提升

指标	分离操作	融合	改进
内存操作	6	2	3 倍减少
内核启动	3	1	3 倍减少
内存流量	高	低	显著

使用

typescript

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

const fused = new Conv2dBiasReLUOperator(context);

const output = await fused.forward([input, weight, bias], {
  channels: 32,
  kernelSize: 3,
  stride: 1,
  padding: 1
});

WGSL 实现

与分离卷积的关键区别：

wgsl

// ... 卷积计算 ...
var sum = 0.0;
// [卷积循环]

// 在同一内核中添加偏置和 ReLU
sum = sum + bias[oc];
sum = select(0.0, sum, sum > 0.0);  // ReLU

output[outIdx] = sum;

何时使用

✅ 对以下情况使用融合算子：

卷积→偏置→ReLU 序列
生产推理
内存受限环境

❌ 不要在以下情况使用：

需要中间卷积输出
调试单独操作
探索不同激活函数

FlattenOperator

重塑张量，保留批次维度。

描述

展平除第一个（批次）维度外的所有维度。

变换

输入:  [N, C, H, W]
输出: [N, C*H*W]

示例：
[1, 3, 224, 224] → [1, 150528]

实现细节

使用 零拷贝视图 - 无 GPU 内存移动：

typescript

// 内部使用 reshape()
const flat = input.reshape([N, C * H * W]);

使用

typescript

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

const flatten = new FlattenOperator(context);

// 为全连接层展平卷积输出
const output = await flatten.forward([convOutput]);
// 形状：[批次, 通道*高度*宽度]

性能

内存：无额外分配
计算：无 GPU 计算
时间：O(1) - 瞬时

DenseOperator

全连接（密集）层。

描述

带学习权重和可选偏置的矩阵乘法。

公式

输出 = 输入 @ 权重.T + 偏置

其中：
- 输入：[N, in_features]
- 权重：[out_features, in_features]
- 偏置：[out_features]
- 输出：[N, out_features]

参数

名称	类型	默认值	描述
`units`	`number`	必需	输出单元数

使用

typescript

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

const dense = new DenseOperator(context);

const output = await dense.forward([input, weight, bias], {
  units: 128
});

WGSL 实现

wgsl

for (var n = 0u; n < N; n = n + 1u) {
  for (var outIdx = 0u; outIdx < outFeatures; outIdx = outIdx + 1u) {
    var sum = 0.0;
    
    // 点积
    for (var inIdx = 0u; inIdx < inFeatures; inIdx = inIdx + 1u) {
      let inputVal = input[n * inFeatures + inIdx];
      let weightVal = weight[outIdx * inFeatures + inIdx];
      sum = sum + inputVal * weightVal;
    }
    
    // 添加偏置并存储
    output[n * outFeatures + outIdx] = sum + bias[outIdx];
  }
}

性能提示

大权重矩阵受益于工作组级分块
考虑生产推理的量化
批处理多个输入以获得更好的 GPU 利用率

算子对比

算子	FLOPs	内存	典型用途
ReLU	O(N)	O(N)	激活
Softmax	O(N)	O(N)	最终分类
MaxPool	O(N×k²)	O(N)	下采样
Conv2d	O(N×C×K×k²)	O(N×K)	特征提取
Conv2dBiasReLU	O(N×C×K×k²)	O(N×K)	卷积层
Flatten	O(1)	O(1)	形状变换
Dense	O(N×I×O)	O(N×O)	分类

图例： N = 批次, C = 通道, K = 输出通道, k = 卷积核大小, I = 输入特征, O = 输出特征

自定义算子

自定义算子模板

typescript

import { Operator, OperatorParams, Tensor, TensorShape } from 'tiny-dl-inference';

class CustomOperator extends Operator {
  protected compileShader(): string {
    return `
      @group(0) @binding(0) var<storage, read_write> output: array<f32>;
      @group(0) @binding(1) var<storage, read> input: array<f32>;
      
      @compute @workgroup_size(256)
      fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
        let idx = global_id.x;
        if (idx >= arrayLength(&input)) { return; }
        
        // 你的计算
        output[idx] = input[idx] * 2.0;
      }
    `;
  }
  
  computeOutputShape(inputShape: TensorShape): TensorShape {
    // 返回输出形状
    return inputShape;
  }
  
  async forward(inputs: Tensor[], params?: OperatorParams): Promise<Tensor> {
    const input = inputs[0];
    const outputShape = this.computeOutputShape(input.shape);
    const output = new Tensor(this.context, outputShape, { layout: input.layout });
    
    this.ensureInitialized();
    
    const encoder = this.context.createCommandEncoder();
    const pass = encoder.beginComputePass();
    pass.setPipeline(this.pipeline!);
    
    const bindGroup = this.context.createBindGroup({
      layout: this.bindGroupLayout!,
      entries: [
        { binding: 0, resource: { buffer: output.buffer } },
        { binding: 1, resource: { buffer: input.buffer } }
      ]
    });
    
    pass.setBindGroup(0, bindGroup);
    pass.dispatchWorkgroups(Math.ceil(output.size / 256));
    pass.end();
    
    this.context.submit([encoder.finish()]);
    
    return output;
  }
}

最佳实践

使用适当的工作组大小（256 是良好的默认值）
检查边界 以避免越界访问
最小化工作组内的分歧
尽可能重用缓冲区
分析你的自定义算子

算子实现 ​

算子概览 ​

ReLUOperator ​

描述 ​

公式 ​

使用 ​

WGSL 实现 ​

性能 ​

SoftmaxOperator ​

描述 ​

公式 ​

参数 ​

使用 ​

WGSL 实现 ​

输出属性 ​

MaxPoolOperator ​

描述 ​

参数 ​

输出形状 ​

使用 ​

WGSL 实现 ​

性能 ​

Conv2dOperator ​

描述 ​

参数 ​

输入/输出形状 ​

使用 ​

WGSL 实现 ​

性能特征 ​

Conv2dBiasReLUOperator ​

描述 ​

公式 ​

性能提升 ​

使用 ​

WGSL 实现 ​

何时使用 ​

FlattenOperator ​

描述 ​

变换 ​

实现细节 ​

使用 ​

性能 ​

DenseOperator ​

描述 ​

公式 ​

参数 ​

使用 ​

WGSL 实现 ​

性能提示 ​

算子对比 ​

自定义算子 ​

自定义算子模板 ​

最佳实践 ​

算子实现

算子概览

ReLUOperator

描述

公式

使用

WGSL 实现

性能

SoftmaxOperator

描述

公式

参数

使用

WGSL 实现

输出属性

MaxPoolOperator

描述

参数

输出形状

使用

WGSL 实现

性能

Conv2dOperator

描述

参数

输入/输出形状

使用

WGSL 实现

性能特征

Conv2dBiasReLUOperator

描述

公式

性能提升

使用

WGSL 实现

何时使用

FlattenOperator

描述

变换

实现细节

使用

性能

DenseOperator

描述

公式

参数

使用

WGSL 实现

性能提示

算子对比

自定义算子

自定义算子模板

最佳实践