Skip to content
目录

Tensor、后处理与自定义 ONNX 算子

1. Tensor 是对 ORT Tensor 的包装

utils/tensor.jsTensor 类:

  • 内部持有 ort_tensor(来自 onnxruntime-common / backends/onnx.js 导出);
  • 支持从现有 ORT Tensor 构造,或用 (dtype, data, dims) 新建;
  • 使用 Proxy 实现类 NumPy 的整数索引访问(_getitem)。
25:99:packages/transformers/src/utils/tensor.js
export class Tensor {
    ...
    get dims() {
        return this.ort_tensor.dims;
    }
    ...
    constructor(...args) {
        if (isONNXTensor(args[0])) {
            this.ort_tensor = /** @type {ONNXTensor} */ (args[0]);
        } else {
            this.ort_tensor = new ONNXTensor(
                /** @type {DataType} */ (args[0]),
                /** @type {Exclude<import('./maths.js').AnyTypedArray, Uint8ClampedArray>} */ (args[1]),
                args[2],
            );
        }

        return new Proxy(this, {
            get: (obj, key) => {
                if (typeof key === 'string') {
                    let index = Number(key);
                    if (Number.isInteger(index)) {
                        return obj._getitem(index);
                    }
                }

原因:ORT 的 Tensor 偏底层;库需要 maths.js 联动、统一 dtype 语义、以及后续 GPU buffer 生命周期管理locationdispose)。

2. sessionRun 与 Tensor 替换

models/session.jsreplaceTensors 在 ORT 输出上递归把 ORT Tensor 换成自定义 Tensor,以便上层使用扩展方法。

3. TensorOpRegistry:微型 ONNX 算子库

ops/registry.js极小的 ONNX 图(字节数组硬编码)包装成惰性初始化的 session,用于 Resize 等 在预处理/后处理中需要、但想用 ORT 一致执行路径的场景。

34:39:packages/transformers/src/ops/registry.js
export class TensorOpRegistry {
    static session_options = {
        // TODO: Allow for multiple execution providers
        // executionProviders: ['webgpu'],
    };

要点:这是 「用 ORT 跑子图算子」 的技巧;TODO 也表明未来可能让这类算子走 WebGPU EP。

4. Tokenizer / Processor

  • Tokenizer@huggingface/tokenizers + 各模型 tokenization_*.js
  • Processor:组合 feature extractor、image processor 等,对齐 Python Processor 行为。

Pipeline 里 预处理 → 模型 → 后处理 的边界清晰,便于单测(tests/pipelines)。

5. 实践建议

  • 大量元素级运算优先复用 utils/maths.js / Tensor 已有方法,避免重复实现导致数值不一致;
  • 需要 GPU 张量时,注意 webInferenceChain 串行 的交互,避免在链外并发 run

下一篇:性能优化与工程最佳实践

MIT License.