ONNX API 参考

ONNX 集成类和辅助工具的完整 API 文档。

目录

• ONNXManager
• ONNXGenerator
• OnnxGpuIO
• PrecisionDetector
• Precision Types (精度类型)
• ONNXModelTester
• ONNXTestUtils
• Utility Functions (实用函数)

---

ONNXManager

协调 ONNX 模型的加载、推理和资源管理。

路径: src/app/managers/onnx-manager.ts

constructor(modelManager: ModelManager)

创建一个作用域为 ModelManager 的管理器。

方法

loadONNXModel(device, modelPath, cameraMatrix, projectionMatrix, name?, options?): Promise

从 URL 加载 ONNX 模型并创建 DynamicPointCloud。

参数:

• device: GPUDevice - 共享的 WebGPU 设备
• modelPath: string - ONNX 模型的 URL 字符串
• cameraMatrix: Float32Array - 初始视图矩阵 (4×4)
• projectionMatrix: Float32Array - 初始投影矩阵 (4×4)
• name?: string - 可选的模型名称（如未提供则自动生成）
• options?: ONNXLoadOptions - 加载选项

返回: Promise<ModelEntry> - 注册到 ModelManager 的模型条目

ONNXLoadOptions:

interface ONNXLoadOptions {
  staticInference?: boolean;        // 如果为 true，则运行一次推理；如果为 false，则启用每帧更新
  maxPoints?: number;               // 手动覆盖缓冲区分配大小（如未提供则自动检测）
  debugLogging?: boolean;           // 启用调试日志
  precisionConfig?: PrecisionConfig; // 手动精度覆盖配置
}

loadONNXFromFile(device, file, cameraMatrix?, projectionMatrix?): Promise

从浏览器 File 对象加载。可选矩阵默认为单位矩阵（相机位于 (0,0,5) 看向原点）。

参数:

• device: GPUDevice - 共享的 WebGPU 设备
• file: File - 浏览器文件对象
• cameraMatrix?: Float32Array | null - 可选的初始视图矩阵
• projectionMatrix?: Float32Array | null - 可选的初始投影矩阵

返回: Promise<ModelEntry>

updateCameraMatrices(modelName, cameraMatrix, projectionMatrix): Promise

更新指定 ONNX 模型的视图/投影矩阵。注意：目前是一个占位符；动态模型通过 AnimationManager 自动更新。

disposeModel(modelId: string): void

通过 ID 销毁特定的 ONNX 模型。清理 GPU 缓冲区和 ORT 会话。

dispose(): void

销毁所有 ONNX 资源。清理所有生成器和点云。

getGenerator(modelId: string): ONNXGenerator | undefined

获取特定模型的生成器（用于调试/高级用途）。

getPointCloud(modelId: string): DynamicPointCloud | undefined

获取特定模型的点云（用于调试/高级用途）。

hasONNXModels(): boolean

如果管理了任何 ONNX 模型，则返回 true。

getONNXModels(): string[]

返回所有已加载的 ONNX 模型 ID 数组。

getONNXPerformanceStats():

返回性能统计信息： - modelCount: 已加载的 ONNX 模型数量 - totalGenerators: ONNXGenerator 实例总数 - totalPointClouds: DynamicPointCloud 实例总数

---

ONNXGenerator

运行 ONNX 推理并暴露 GPU 缓冲区。提供 OnnxGpuIO 的简化外观（facade）。

路径: src/onnx/onnx_generator.ts

constructor(cfg: ONNXGeneratorConfig)

配置:

interface ONNXGeneratorConfig {
  modelUrl: string;                    // ONNX 模型 URL 或路径
  maxPoints?: number;                  // 可选，将从模型元数据中自动检测
  debugLogging?: boolean;              // 启用调试日志
  device?: GPUDevice;                  // 传递应用的 WebGPU 设备以避免不匹配
  precisionConfig?: PrecisionConfig;   // 手动精度覆盖配置
}

方法

生命周期:

• initialize(device?: GPUDevice): Promise<void> – 初始化生成器。使用 device 参数或 cfg.device。如果无可用的设备则抛出错误。

推理:

• generate(inputData?: { cameraMatrix?: Float32Array, projectionMatrix?: Float32Array, time?: number }): Promise<void> – 使用可选输入运行推理。对于静态模型，可以用 {} 调用。对于动态模型，提供相机矩阵和时间。

GPU 缓冲区访问:

• getGaussianBuffer(): GPUBuffer – 获取预分配的高斯数据缓冲区
• getSHBuffer(): GPUBuffer – 获取预分配的颜色 (SH/RGB) 数据缓冲区
• getCountBuffer(): GPUBuffer – 获取点计数缓冲区 (int32)

设备与输入:

• getDevice(): GPUDevice – 获取 WebGPU 设备
• getInputNames(): readonly string[] – 获取模型预期的输入名称

元数据:

• getDetectedCapacity(): number – 从模型元数据获取检测到的最大点数
• getDetectedColorMode(): 'sh' | 'rgb' – 获取检测到的颜色模式
• getDetectedColorDim(): number – 获取检测到的颜色维度（SH 为 48，RGB 为 3）
• getActualMaxPoints(): number – 获取实际最大点数（检测到的或配置的）

精度信息:

• getGaussianPrecision(): PrecisionMetadata – 获取高斯输出的精度元数据
• getColorPrecision(): PrecisionMetadata – 获取颜色输出的精度元数据

清理:

• dispose(): void – 释放所有 GPU 资源和 ORT 会话

---

OnnxGpuIO

处理与 ONNX Runtime 的底层 GPU I/O 绑定。管理会话创建、缓冲区分配和推理执行。

路径: src/onnx/onnx_gpu_io.ts

constructor()

创建一个新的 OnnxGpuIO 实例。使用前必须调用 init()。

方法

init(cfg: OnnxGpuIOConfig & { precisionConfig?: PrecisionConfig }): Promise

使用配置初始化 I/O。

配置:

interface OnnxGpuIOConfig {
  modelUrl: string;            // ONNX 模型路径
  maxPoints?: number;          // 预设最大点数（可选，从元数据推断）
  device: GPUDevice;          // 使用应用现有的 WebGPU 设备
  verbose?: boolean;           // 启用详细调试日志
  precisionConfig?: PrecisionConfig; // 手动精度覆盖配置
}

流程:

1. 初始化 ONNX Runtime 环境
2. 创建仅 WebGPU 的 InferenceSession（带有图捕获回退）
3. 从元数据检测容量和颜色模式
4. 检测精度（或应用手动配置）
5. 预分配所有 GPU 缓冲区，具有适当的大小和对齐

updateInputBuffers(view?: mat4, proj?: mat4, time?: number): void

将输入写入预分配的 GPU 缓冲区。更新 cameraMatrixBuf、projMatrixBuf 和 timeBuf。

runInference(input?: { cameraMatrix?: Float32Array, projectionMatrix?: Float32Array, time?: number }): Promise

执行一次推理。使用独占执行链来防止并发冲突。将 GPU 缓冲区绑定为 feeds/fetches 并执行会话，无需 CPU 往返。

注意: 内部使用 OnnxGpuIO.runExclusive() 来序列化推理调用。

destroy(): void

释放所有资源：ORT 会话、GPU 缓冲区和设备引用。

静态方法

runExclusive(fn: () => Promise): Promise

全局独占执行协调器。确保一次只运行一个推理，以防止 ORT WebGPU IOBinding 会话冲突。

公共属性

输出缓冲区:

• gaussBuf: GPUBuffer – 预分配的高斯数据缓冲区
• shBuf: GPUBuffer – 预分配的颜色 (SH/RGB) 数据缓冲区
• countBuf: GPUBuffer – 点计数缓冲区 (int32)

输入缓冲区:

• cameraMatrixBuf: GPUBuffer – 相机视图矩阵缓冲区 (4×4 float32)
• projMatrixBuf: GPUBuffer – 投影矩阵缓冲区 (4×4 float32)
• timeBuf: GPUBuffer – 时间输入缓冲区 (float32)

会话与设备:

• session: ort.InferenceSession – ONNX Runtime 推理会话
• device: GPUDevice – WebGPU 设备

元数据:

• inputNames: readonly string[] – 模型预期的输入名称
• maxPoints: number – 实际最大点数（检测到的或配置的）
• actualPoints: number – 模型返回的实际点数

检测结果:

• detectedCapacity: number – 从元数据检测到的最大点数
• detectedColorMode: 'sh' | 'rgb' – 检测到的颜色模式
• detectedColorDim: number – 检测到的颜色维度
• detectedGaussOutputName: string | null – 检测到的高斯输出名称
• detectedGaussFields: number – 高斯字段的数量（通常为 10）
• detectedColorOutputName: string | null – 检测到的颜色输出名称

---

PrecisionDetector

针对 ONNX 模型输出的自动精度检测。

路径: src/onnx/precision-detector.ts

静态方法

detectOutputPrecisionFromName(outputName: string): PrecisionMetadata

根据输出名称后缀检测精度。

后缀模式:

• _f32, _float32 → float32 (4 bytes)
• _f16, _float16 → float16 (2 bytes)
• _i8, _int8 → int8 (1 byte)
• _u8, _uint8 → uint8 (1 byte)
• 默认 → float16 (2 bytes)

detectFromMetadataPreferringNameSuffix(session: ort.InferenceSession, outputName: string): PrecisionMetadata

检测精度的优先级： 1. 检查会话输出元数据中的类型信息 2. 如果元数据不可用，则回退到基于名称的检测

extractQuantizationParams(session: ort.InferenceSession, tensorName: string):

从模型初始化器（initializers）中提取量化参数（scale, zeroPoint）。从模型图中尽力提取。

calculateBufferSize(dims: number[], precision: PrecisionMetadata): number

计算缓冲区大小，并进行 16 字节对齐。

---

Precision Types (精度类型)

精度处理的类型定义。

路径: src/onnx/precision-types.ts

类型

type OnnxDataType = 'float32' | 'float16' | 'int8' | 'uint8';

interface PrecisionMetadata {
  dataType: OnnxDataType;
  bytesPerElement: number;
  scale?: number;      // 用于量化的 int8/uint8
  zeroPoint?: number;  // 用于量化的 int8/uint8
}

interface OutputBufferDescriptor {
  name: string;
  precision: PrecisionMetadata;
  dims: number[];
  sizeInBytes: number;
}

interface PrecisionConfig {
  gaussian?: Partial<PrecisionMetadata>;  // 高斯输出的覆盖配置
  color?: Partial<PrecisionMetadata>;     // 颜色输出的覆盖配置
  autoDetect?: boolean;                    // 遗留标志（已弃用）
}

实用函数

• align16(n: number): number – 将数字对齐到 16 字节边界
• calcSizeInBytes(dims: number[], p: PrecisionMetadata): number – 计算缓冲区大小（含对齐）
• dataTypeToOrtString(p: PrecisionMetadata): 'float16' | 'float32' | 'int8' | 'uint8' – 转换为 ORT 类型字符串

---

ONNXModelTester

用于隔离加载和验证 ONNX 模型的实用程序。

静态方法

• initialize(): Promise<void> – 设置 ONNX Runtime。
• loadModel(modelPath?: string): Promise<void> – 默认路径 './models/gaussians3d.onnx'。
• testInference(inputs?: any): Promise<any> – 运行测试推理。
• dispose(): void

---

ONNXTestUtils

用于调试和基准测试的助手。

静态方法

• printInferenceReport(generator): void
• compareWithReference(generator, referenceData): void
• validateOutputBuffers(generator): boolean
• measurePerformance(generator, iterations = 100): PerformanceMetrics

interface PerformanceMetrics {
  averageTime: number;
  minTime: number;
  maxTime: number;
  totalTime: number;
  iterations: number;
}

---

Utility Functions (实用函数)

• testONNXModel(modelUrl: string, device: GPUDevice): Promise<TestResult>
• runONNXIntegrationTest(): Promise<void>
• runONNXPerformanceTest(): Promise<PerformanceMetrics>

interface TestResult {
  success: boolean;
  error?: string;
  performance?: PerformanceMetrics;
  outputValidation?: boolean;
}

---

使用示例

基本 ONNX Generator 用法

import { ONNXGenerator } from './onnx/onnx_generator';

const generator = new ONNXGenerator({
  modelUrl: '/models/gaussians3d.onnx',
  maxPoints: 1_000_000,
  debugLogging: true,
  device: gpuDevice
});

await generator.initialize();
await generator.generate({
  cameraMatrix: viewMatrix,
  projectionMatrix: projMatrix,
  time: performance.now() / 1000,
});

const gaussianBuffer = generator.getGaussianBuffer();
const shBuffer = generator.getSHBuffer();
const countBuffer = generator.getCountBuffer();

// 访问精度信息
const gaussPrecision = generator.getGaussianPrecision();
const colorPrecision = generator.getColorPrecision();
console.log(`Gaussian: ${gaussPrecision.dataType}, Color: ${colorPrecision.dataType}`);

使用 ONNXManager

import { ONNXManager } from './app/managers/onnx-manager';

const onnxManager = new ONNXManager(modelManager);

// 加载静态模型
const staticEntry = await onnxManager.loadONNXModel(
  device,
  '/models/static.onnx',
  cameraMatrix,
  projectionMatrix,
  'static-model',
  { 
    staticInference: true,
    maxPoints: 2_000_000,
    debugLogging: true
  }
);

// 加载动态模型（每帧更新）
const dynamicEntry = await onnxManager.loadONNXModel(
  device,
  '/models/dynamic.onnx',
  cameraMatrix,
  projectionMatrix,
  'dynamic-model',
  { 
    staticInference: false,  // 启用每帧更新
    maxPoints: 2_000_000,
    debugLogging: true
  }
);

// 使用精度覆盖加载
const quantizedEntry = await onnxManager.loadONNXModel(
  device,
  '/models/quantized.onnx',
  cameraMatrix,
  projectionMatrix,
  'quantized-model',
  {
    precisionConfig: {
      gaussian: { dataType: 'int8', bytesPerElement: 1 },
      color: { dataType: 'int8', bytesPerElement: 1 }
    }
  }
);

// 获取生成器用于高级用途
const generator = onnxManager.getGenerator(dynamicEntry.id);
if (generator) {
  console.log(`Color mode: ${generator.getDetectedColorMode()}`);
  console.log(`Capacity: ${generator.getDetectedCapacity()}`);
}

// 清理
onnxManager.disposeModel(staticEntry.id);
// 或者销毁所有
onnxManager.dispose();

性能测试

import { ONNXTestUtils } from './test_utils';

const metrics = ONNXTestUtils.measurePerformance(generator, 1000);
console.log(`Average: ${metrics.averageTime} ms`);
console.log(`Min: ${metrics.minTime} ms`);
console.log(`Max: ${metrics.maxTime} ms`);

ONNXTestUtils.printInferenceReport(generator);

模型测试

import { testONNXModel } from './test_loader';

const result = await testONNXModel('/models/test.onnx', device);
if (result.success) {
  console.log('Model test passed');
  console.log('Performance:', result.performance);
} else {
  console.error('Model test failed:', result.error);
}

---

注意事项

1. 共享设备 – ONNX 管道重用应用的 GPUDevice，保证缓冲区兼容性并避免设备不匹配错误。

2. 仅限 GPU – 无 CPU 往返；所有缓冲区都留在 GPU 上。输入和输出直接作为 GPU 缓冲区绑定。

3. 动态模型 – 支持每帧更新而无需重新分配资源。DynamicPointCloud 与 ONNXGenerator 连接以实现自动更新。

4. 精度检测 – 从模型元数据或输出名称自动检测数据类型。支持通过 PrecisionConfig 进行手动覆盖。

5. 独占执行 – 使用全局执行链来防止与 ORT WebGPU IOBinding 的并发推理冲突。

6. 图捕获 – 支持 WebGPU 图捕获以提高性能（如果不支持则自动回退）。

7. 性能 – 设备共享和缓冲区重用保持推理快速。预分配缓冲区消除了分配开销。

8. 调试 – 辅助工具（ONNXTestUtils, ONNXModelTester）简化了验证和分析。

9. 兼容性 – 适用于静态和动态推理，具有自动元数据检测（容量、颜色模式、精度）。

10. 资源管理 – 通过 dispose() 方法进行适当的清理，确保释放 GPU 资源。