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. Suy luận tăng tốc phần cứng trong trình duyệt: Chạy các mô hình có hiệu suất CPU vượt trội, được tăng tốc bằng XNNPack được liên kết với WebAssembly (Wasm) gọn nhẹ. Đối với GPU và hoạt động mở rộng quy mô phần cứng chuyên dụng (chẳng hạn như NPU), LiteRT.js tự nhiên hiển thị cả API WebGPU và API WebNN mới nổi, giúp tối ưu hoá chi tiết theo từng nền tảng cụ thể.
  2. Khả năng tương thích với nhiều khung: Đơn giản hoá ngữ nghĩa phát triển bằng cách biên dịch từ Khung học máy mà bạn muốn một cách tự nhiên: PyTorch, JAX hoặc TensorFlow.
  3. Lặp lại các quy trình hiện có: Tích hợp sẵn với các cấu trúc TensorFlow.js hiện có bằng cách phân tích cú pháp các Tensor TensorFlow.js được hỗ trợ gốc dưới dạng đầu vào và đầu ra ranh giới trực tiếp.

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 đó, nhập gói và tải các tệp Wasm:

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

// Load the LiteRT.js Wasm files from a CDN.
await loadLiteRt('https://cdn.jsdelivr.net/npm/@litertjs/core/wasm/')
// Alternatively, host them from your server.
// They are located in node_modules/@litertjs/core/wasm/
await loadLiteRt(`your/path/to/wasm/`);

Chuyển đổi mô hình

LiteRT.js sử dụng cùng định dạng .tflite như phần còn lại của hệ sinh thái LiteRT và hỗ trợ các mô hình hiện có trên Kaggle và Huggingface. Nếu có mô hình PyTorch mới, bạn cần chuyển đổi mô hình đó.

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

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

import litert_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 = litert_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',
    // Can select from 'webnn', 'webgpu', & 'wasm'.
    // Additionally, you can pass an array of accelerators e.g. ['webnn', 'wasm']
    // if you would like to fallback to CPU execution,
    // Note that ONLY cpu fallback is supported for now
    // (i.e. specifying ['webnn', 'webgpu']) will lead to compilation errors
});
// 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 = new Tensor(image, /* shape */ [1, 3, 224, 224]);

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

// Clean up and get outputs
inputTensor.delete();
const output = outputs[0];
const outputData = await output.data();
output.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 GPU và phần cứng vượt trội: Các mô hình LiteRT.js tận dụng khả năng tăng tốc WebGPU để tối ưu hoá hiệu suất trên các cấu trúc trình duyệt. Với sự hỗ trợ cho WebGPU và WebNN sắp ra mắt, LiteRT.js mang đến khả năng tăng tốc phần cứng linh hoạt trên nhiều thiết bị biên.
  2. Lộ trình chuyển đổi mô hình dễ dàng hơn: Lộ trình chuyển đổi LiteRT.js đi trực tiếp từ PyTorch đến LiteRT. Đường dẫn chuyển đổi từ PyTorch sang TensorFlow.js phức tạp hơn đáng kể, yêu cầu bạn 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, thứ duy nhất bạn cần di chuyển là 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 sang .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 NPM @litertjs/core và @litertjs/tfjs-interop.
  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ế tải mô hình TensorFlow.js bằng tải mô hình LiteRT.js.
  5. Thay thế model.predict(inputs) hoặc model.execute(inputs) của TensorFlow.js bằng runWithTfjsTensors(liteRtModel, inputs). runWithTfjsTensors lấy cùng các tensor đầu vào mà các mô hình TensorFlow.js sử dụng và xuất ra 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 dữ liệu đầ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().