Sapphire开发日志（十二）——技术重难点2：用 ONNX-Runtime-Web 运行 ONNX 推理的封装与使用

用 ONNX-Runtime-Web 运行 ONNX 推理的封装与使用

ONNX Runtime Web 是一个高性能的推理引擎，支持在浏览器中运行 ONNX 模型。通过封装 ONNX Runtime Web，可以简化前端项目中模型推理的集成过程，提高代码的可维护性和可复用性。

安装和配置

首先，我们需要安装 onnxruntime-web 依赖。

npm install onnxruntime-web

创建封装模块

为了简化 ONNX 模型的加载和推理过程，我们可以创建一个封装模块。例如，我们可以创建一个 useOnnxModel 钩子。

import { InferenceSession } from 'onnxruntime-web';
import { useState, useEffect } from 'react';

const useOnnxModel = (modelUrl) => {
  const [session, setSession] = useState(null);
  const [isLoading, setIsLoading] = useState(true);
  const [error, setError] = useState(null);

  useEffect(() => {
    const loadModel = async () => {
      try {
        const session = await InferenceSession.create(modelUrl);
        setSession(session);
      } catch (err) {
        setError(err);
      } finally {
        setIsLoading(false);
      }
    };

    loadModel();
  }, [modelUrl]);

  const runInference = async (inputData) => {
    if (!session) {
      throw new Error('Model is not loaded yet');
    }
    const feeds = { input: new Float32Array(inputData) };
    const output = await session.run(feeds);
    return output;
  };

  return { session, isLoading, error, runInference };
};

export default useOnnxModel;

在项目中接入

现在，我们可以在项目中使用我们封装的 useOnnxModel 钩子来进行模型推理。

import React, { useState } from 'react';
import useOnnxModel from './hooks/useOnnxModel';

const ModelInferenceComponent = () => {
  const modelUrl = 'path/to/your/model.onnx';
  const { session, isLoading, error, runInference } = useOnnxModel(modelUrl);
  const [inputData, setInputData] = useState([/* Your input data here */]);
  const [outputData, setOutputData] = useState(null);

  const handleInference = async () => {
    try {
      const output = await runInference(inputData);
      setOutputData(output);
    } catch (err) {
      console.error('Inference error:', err);
    }
  };

  if (isLoading) return <div>Loading model...</div>;
  if (error) return <div>Error loading model: {error.message}</div>;

  return (
    <div>
      <h1>ONNX Model Inference</h1>
      <button onClick={handleInference}>Run Inference</button>
      {outputData && <div>Output: {JSON.stringify(outputData)}</div>}
    </div>
  );
};

export default ModelInferenceComponent;

实践应用

以下是一个应用场景，包括如何在前端项目中加载并运行 ONNX 模型。

// hooks/useOnnxModel.js
import { InferenceSession } from 'onnxruntime-web';
import { useState, useEffect } from 'react';

const useOnnxModel = (modelUrl) => {
  const [session, setSession] = useState(null);
  const [isLoading, setIsLoading] = useState(true);
  const [error, setError] = useState(null);

  useEffect(() => {
    const loadModel = async () => {
      try {
        const session = await InferenceSession.create(modelUrl);
        setSession(session);
      } catch (err) {
        setError(err);
      } finally {
        setIsLoading(false);
      }
    };

    loadModel();
  }, [modelUrl]);

  const runInference = async (inputData) => {
    if (!session) {
      throw new Error('Model is not loaded yet');
    }
    const feeds = { input: new Float32Array(inputData) };
    const output = await session.run(feeds);
    return output;
  };

  return { session, isLoading, error, runInference };
};

export default useOnnxModel;

// components/ModelInferenceComponent.js
import React, { useState } from 'react';
import useOnnxModel from '../hooks/useOnnxModel';

const ModelInferenceComponent = () => {
  const modelUrl = 'path/to/your/model.onnx';
  const { session, isLoading, error, runInference } = useOnnxModel(modelUrl);
  const [inputData, setInputData] = useState([/* Your input data here */]);
  const [outputData, setOutputData] = useState(null);

  const handleInference = async () => {
    try {
      const output = await runInference(inputData);
      setOutputData(output);
    } catch (err) {
      console.error('Inference error:', err);
    }
  };

  if (isLoading) return <div>Loading model...</div>;
  if (error) return <div>Error loading model: {error.message}</div>;

  return (
    <div>
      <h1>ONNX Model Inference</h1>
      <button onClick={handleInference}>Run Inference</button>
      {outputData && <div>Output: {JSON.stringify(outputData)}</div>}
    </div>
  );
};

export default ModelInferenceComponent;

// App.js
import React from 'react';
import ModelInferenceComponent from './components/ModelInferenceComponent';

const App = () => (
  <div>
    <h1>ONNX Runtime Web Example</h1>
    <ModelInferenceComponent />
  </div>
);

export default App;

通过封装 ONNX Runtime Web，我们可以简化前端项目中模型推理的集成过程，提高代码的可维护性和可复用性。在项目中接入 ONNX Runtime Web 后，我们可以在浏览器中高效地运行 ONNX 模型，提供更好的用户体验。