LiteRT cho web bằng LiteRT.js

LiteRT.js là thời gian chạy WebAI hiệu suất cao của Google, nhắm đến các ứng dụng Web sản xuất. Đây là phần tiếp nối của ngăn xếp LiteRT, đảm bảo khả năng hỗ trợ đa khung và hợp nhất thời gian chạy cốt lõi của chúng tôi trên tất cả các nền tảng.

LiteRT.js hỗ trợ các tính năng cốt lõi sau:

  1. Hỗ trợ mô hình LiteRT trong trình duyệt: Chạy các mô hình có hiệu suất tốt nhất trên CPU, được tăng tốc thông qua XNNPack trên WebAssembly (Wasm) và GPU bằng API WebGPU.
  2. Khả năng tương thích với nhiều khung: Sử dụng Khung máy học mà bạn muốn: PyTorch, Jax hoặc TensorFlow.
  3. Xây dựng trên các quy trình hiện có: Tích hợp với các quy trình TensorFlow.js hiện có bằng cách hỗ trợ Tensor TensorFlow.js làm đầu vào và đầu ra.

Cài đặt

Cài đặt gói @litertjs/core từ npm:

npm install @litertjs/core

Các tệp Wasm nằm trong node_modules/@litertjs/core/wasm/. Để thuận tiện, hãy sao chép và phân phát toàn bộ thư mục wasm/. Sau đó, hãy nhập gói và tải các tệp Wasm:

import {loadLiteRt} from '@litertjs/core;

// Host LiteRT's Wasm files on your server.
await loadLiteRt(`your/path/to/wasm/`);

Chuyển đổi mô hình

LiteRT.js sử dụng cùng định dạng .tflite như Android và iOS, đồng thời hỗ trợ các mô hình hiện có trên Kaggle và Huggingface. Nếu có một mô hình PyTorch mới, bạn sẽ cần chuyển đổi mô hình đó.

Chuyển đổi mô hình PyTorch sang LiteRT

Để chuyển đổi mô hình PyTorch sang LiteRT, hãy sử dụng ai-edge-torch converter.

import ai_edge_torch

# Load your torch model. We're using resnet for this example.
resnet18 = torchvision.models.resnet18(torchvision.models.ResNet18_Weights.IMAGENET1K_V1)

sample_inputs = (torch.randn(1, 3, 224, 224),)

# Convert the model to LiteRT.
edge_model = ai_edge_torch.convert(resnet18.eval(), sample_inputs)

# Export the model.
edge_model.export('resnet.tflite')

Chạy Mô hình đã chuyển đổi

Sau khi chuyển đổi mô hình thành tệp .tflite, bạn có thể chạy mô hình đó trong trình duyệt.

import {loadAndCompile} from '@litertjs/core';

// Load the model hosted from your server. This makes an http(s) request.
const model = await loadAndCompile('/path/to/model.tflite', {
    accelerator: 'webgpu', // or 'wasm' for XNNPack CPU inference
});
// The model can also be loaded from a Uint8Array if you want to fetch it yourself.

// Create image input data
const image = new Float32Array(224 * 224 * 3).fill(0);
const inputTensor =
    await new Tensor(image, /* shape */ [1, 3, 224, 224]).moveTo('webgpu');

// Run the model
const outputs = model(inputTensor);
// You can also use model([inputTensor])
// or model({'input_tensor_name': inputTensor})

// Clean up and get outputs
inputTensor.delete();
const outputTensorCpu = await outputs[0].moveTo('wasm');
const outputData = outputTensorCpu.toTypedArray();
outputTensorCpu.delete();

Tích hợp vào các quy trình TensorFlow.js hiện có

Bạn nên cân nhắc việc tích hợp LiteRT.js vào quy trình TensorFlow.js vì những lý do sau:

  1. Hiệu suất WebGPU tốt nhất: Các mô hình đã chuyển đổi chạy trên WebGPU LiteRT.js được tối ưu hoá cho hiệu suất trình duyệt và đặc biệt nhanh trên các trình duyệt dựa trên Chromium.
  2. Đường dẫn chuyển đổi mô hình dễ dàng hơn: Đường dẫn chuyển đổi LiteRT.js chuyển trực tiếp từ PyTorch sang LiteRT. Đường dẫn chuyển đổi PyTorch sang TensorFlow.js phức tạp hơn đáng kể, đòi hỏi bạn phải chuyển từ PyTorch -> ONNX -> TensorFlow -> TensorFlow.js.
  3. Công cụ gỡ lỗi: Lộ trình chuyển đổi LiteRT.js đi kèm với công cụ gỡ lỗi.

LiteRT.js được thiết kế để hoạt động trong các quy trình TensorFlow.js và tương thích với quy trình xử lý trước và sau của TensorFlow.js, vì vậy, bạn chỉ cần di chuyển chính mô hình.

Tích hợp LiteRT.js vào các quy trình TensorFlow.js theo các bước sau:

  1. Chuyển đổi mô hình TensorFlow, JAX hoặc PyTorch ban đầu của bạn thành .tflite. Để biết thông tin chi tiết, hãy xem phần chuyển đổi mô hình.
  2. Cài đặt các gói @litertjs/core và @litertjs/tfjs-interop NPM.
  3. Nhập và sử dụng phần phụ trợ TensorFlow.js WebGPU. Đây là điều kiện bắt buộc để LiteRT.js tương tác với TensorFlow.js.
  4. Thay thế việc tải mô hình TensorFlow.js bằng việc tải mô hình LiteRT.js.
  5. Thay thế TensorFlow.js model.predict(inputs) hoặc model.execute(inputs) bằng runWithTfjsTensors(liteRtModel, inputs). runWithTfjsTensors lấy các tensor đầu vào mà mô hình TensorFlow.js sử dụng và xuất các tensor TensorFlow.js.
  6. Kiểm thử để đảm bảo quy trình mô hình xuất ra kết quả như bạn mong đợi.

Việc sử dụng LiteRT.js với runWithTfjsTensors cũng có thể yêu cầu bạn thực hiện những thay đổi sau đối với các đầu vào của mô hình:

  1. Sắp xếp lại các đầu vào: Tuỳ thuộc vào cách bộ chuyển đổi sắp xếp các đầu vào và đầu ra của mô hình, bạn có thể cần thay đổi thứ tự của chúng khi truyền các đầu vào và đầu ra đó.
  2. Chuyển vị đầu vào: Cũng có thể trình chuyển đổi đã thay đổi bố cục của đầu vào và đầu ra của mô hình so với những gì TensorFlow.js sử dụng. Bạn có thể cần chuyển vị dữ liệu đầu vào để khớp với mô hình và dữ liệu đầu ra để khớp với phần còn lại của quy trình.
  3. Đổi tên các đầu vào: Nếu bạn đang sử dụng các đầu vào có tên, thì tên của các đầu vào đó cũng có thể đã thay đổi.

Bạn có thể xem thêm thông tin về đầu vào và đầu ra của mô hình bằng model.getInputDetails() và model.getOutputDetails().