Tạo các gói trình duyệt được tối ưu hóa kích thước bằng TensorFlow.js

Tổng quan

TensorFlow.js 3.0 mang đến sự hỗ trợ cho việc xây dựng các gói trình duyệt hướng sản xuất, được tối ưu hóa kích thước . Nói cách khác, chúng tôi muốn giúp bạn gửi ít JavaScript hơn tới trình duyệt một cách dễ dàng hơn.

Tính năng này hướng tới người dùng trong các trường hợp sử dụng sản xuất, những người sẽ đặc biệt được hưởng lợi từ việc loại bỏ byte khỏi tải trọng của họ (và do đó sẵn sàng nỗ lực để đạt được điều này). Để sử dụng tính năng này, bạn nên làm quen với Mô-đun ES , các công cụ đóng gói JavaScript như webpack hoặc rollup và các khái niệm như rung cây/loại bỏ mã chết .

Hướng dẫn này trình bày cách tạo mô-đun tensorflow.js tùy chỉnh có thể được sử dụng với trình đóng gói để tạo bản dựng được tối ưu hóa kích thước cho chương trình bằng tensorflow.js.

Thuật ngữ

Trong ngữ cảnh của tài liệu này, chúng tôi sẽ sử dụng một số thuật ngữ chính:

ES Modules - Hệ thống mô-đun JavaScript tiêu chuẩn . Được giới thiệu trong ES6/ES2015. Có thể nhận dạng bằng cách sử dụng báo cáo nhậpxuất .

Gói - Lấy một tập hợp các nội dung JavaScript và nhóm/gói chúng thành một hoặc nhiều nội dung JavaScript có thể sử dụng được trong trình duyệt. Đây là bước thường tạo ra nội dung cuối cùng được cung cấp cho trình duyệt. Các ứng dụng thường sẽ tự đóng gói trực tiếp từ các nguồn thư viện được dịch mã . Các gói phổ biến bao gồm rollupwebpack . Kết quả cuối cùng của việc gộp gói được gọi là một gói (hoặc đôi khi là một đoạn nếu nó được chia thành nhiều phần)

Rung cây / Loại bỏ mã chết - Loại bỏ mã không được ứng dụng viết cuối cùng sử dụng. Điều này được thực hiện trong quá trình đóng gói, thường là ở bước rút gọn.

Hoạt động (Ops) - Một phép toán trên một hoặc nhiều tensor tạo ra một hoặc nhiều tensor làm đầu ra. Ops là mã 'cấp cao' và có thể sử dụng các ops khác để xác định logic của chúng.

Kernel - Việc triển khai cụ thể một op gắn liền với các khả năng phần cứng cụ thể. Hạt nhân là 'cấp thấp' và phụ trợ cụ thể. Một số op có ánh xạ một-một từ op đến kernel trong khi các op khác sử dụng nhiều kernel.

Phạm vi và trường hợp sử dụng

Chỉ suy luận các mô hình đồ thị

Trường hợp sử dụng chính mà chúng tôi đã nghe được từ những người dùng liên quan đến vấn đề này và đang hỗ trợ trong bản phát hành này là trường hợp thực hiện suy luận với các mô hình biểu đồ TensorFlow.js . Nếu bạn đang sử dụng mô hình lớp TensorFlow.js , bạn có thể chuyển đổi mô hình này sang định dạng mô hình đồ thị bằng cách sử dụng tfjs-converter . Định dạng mô hình đồ thị hiệu quả hơn cho trường hợp sử dụng suy luận.

Thao tác Tensor ở mức độ thấp với tfjs-core

Trường hợp sử dụng khác mà chúng tôi hỗ trợ là các chương trình sử dụng trực tiếp gói @tensorflow/tjfs-core để thao tác với tensor cấp thấp hơn.

Cách tiếp cận của chúng tôi đối với các bản dựng tùy chỉnh

Các nguyên tắc cốt lõi của chúng tôi khi thiết kế chức năng này bao gồm:

  • Tận dụng tối đa hệ thống mô-đun JavaScript (ESM) và cho phép người dùng TensorFlow.js làm điều tương tự.
  • Làm cho TensorFlow.js có khả năng rung cây tốt nhất có thể bằng các gói hiện có (ví dụ: webpack, rollup, v.v.). Điều này cho phép người dùng tận dụng tất cả khả năng của các gói đó, bao gồm các tính năng như tách mã.
  • Duy trì sự dễ sử dụng ở mức tối đa cho những người dùng không quá nhạy cảm với kích thước gói . Điều này không có nghĩa là các bản dựng sản xuất sẽ đòi hỏi nhiều nỗ lực hơn vì nhiều giá trị mặc định trong thư viện của chúng tôi hỗ trợ tính dễ sử dụng so với các bản dựng được tối ưu hóa về kích thước.

Mục tiêu chính trong quy trình làm việc của chúng tôi là tạo ra một mô-đun JavaScript tùy chỉnh cho TensorFlow.js chỉ chứa chức năng cần thiết cho chương trình mà chúng tôi đang cố gắng tối ưu hóa. Chúng tôi dựa vào các gói hiện có để thực hiện tối ưu hóa thực tế.

Mặc dù chúng tôi chủ yếu dựa vào hệ thống mô-đun JavaScript, nhưng chúng tôi cũng cung cấp công cụ CLI tùy chỉnh để xử lý các phần không dễ chỉ định thông qua hệ thống mô-đun trong mã mà người dùng phải đối mặt. Hai ví dụ về điều này là:

  • Thông số kỹ thuật của mô hình được lưu trữ trong tệp model.json
  • Hệ thống điều phối hạt nhân dành riêng cho chương trình phụ trợ mà chúng tôi sử dụng.

Điều này làm cho việc tạo một bản dựng tfjs tùy chỉnh phức tạp hơn một chút so với việc chỉ trỏ một gói vào gói @tensorflow/tfjs thông thường.

Cách tạo gói tùy chỉnh được tối ưu hóa kích thước

Bước 1: Xác định kernel nào chương trình của bạn đang sử dụng

Bước này cho phép chúng tôi xác định tất cả các hạt nhân được sử dụng bởi bất kỳ mô hình nào bạn chạy hoặc mã xử lý trước/sau dựa trên phần phụ trợ mà bạn đã chọn.

Sử dụng tf.profile để chạy các phần ứng dụng sử dụng tensorflow.js và lấy kernel. Nó sẽ trông giống như thế này

const profileInfo = await tf.profile(() => {
  // You must profile all uses of tf symbols.
  runAllMyTfjsCode();
});

const kernelNames = profileInfo.kernelNames
console.log(kernelNames);

Sao chép danh sách hạt nhân đó vào khay nhớ tạm của bạn cho bước tiếp theo.

Bạn cần lập hồ sơ mã bằng cách sử dụng cùng (các) chương trình phụ trợ mà bạn muốn sử dụng trong gói tùy chỉnh của mình.

Bạn sẽ cần lặp lại bước này nếu mô hình của bạn thay đổi hoặc mã xử lý trước/sau của bạn thay đổi.

Bước 2. Viết tệp cấu hình cho mô-đun tfjs tùy chỉnh

Đây là một tập tin cấu hình ví dụ.

Nó trông như thế này:

{
  "kernels": ["Reshape", "_FusedMatMul", "Identity"],
  "backends": [
      "cpu"
  ],
  "models": [
      "./model/model.json"
  ],
  "outputPath": "./custom_tfjs",
  "forwardModeOnly": true
}
  • hạt nhân: Danh sách các hạt nhân sẽ có trong gói. Sao chép phần này từ đầu ra của Bước 1.
  • phụ trợ: Danh sách (các) phụ trợ bạn muốn đưa vào. Các tùy chọn hợp lệ bao gồm "cpu", "webgl" và "wasm".
  • mô hình: Danh sách các tệp model.json cho các mô hình bạn tải trong ứng dụng của mình. Có thể trống nếu chương trình của bạn không sử dụng tfjs_converter để tải mô hình biểu đồ.
  • OutputPath: Đường dẫn đến thư mục để đặt các mô-đun được tạo vào.
  • ForwardModeOnly: Đặt giá trị này thành false nếu bạn muốn bao gồm độ dốc cho các hạt nhân được liệt kê trước đó.

Bước 3. Tạo mô-đun tfjs tùy chỉnh

Chạy công cụ xây dựng tùy chỉnh với tệp cấu hình làm đối số. Bạn cần cài đặt gói @tensorflow/tfjs để có quyền truy cập vào công cụ này.

npx tfjs-custom-module  --config custom_tfjs_config.json

Điều này sẽ tạo một thư mục tại outputPath với một số tệp mới.

Bước 4. Định cấu hình trình đóng gói của bạn thành bí danh tfjs cho mô-đun tùy chỉnh mới.

Trong các gói như webpack và rollup, chúng ta có thể đặt bí danh cho các tham chiếu hiện có cho các mô-đun tfjs để trỏ đến các mô-đun tfjs tùy chỉnh mới được tạo của chúng ta. Có ba mô-đun cần được đặt bí danh để tiết kiệm tối đa kích thước gói.

Đây là một đoạn mã trông như thế nào trong webpack ( ví dụ đầy đủ ở đây ):

...

config.resolve = {
  alias: {
    '@tensorflow/tfjs$':
        path.resolve(__dirname, './custom_tfjs/custom_tfjs.js'),
    '@tensorflow/tfjs-core$': path.resolve(
        __dirname, './custom_tfjs/custom_tfjs_core.js'),
    '@tensorflow/tfjs-core/dist/ops/ops_for_converter': path.resolve(
        __dirname, './custom_tfjs/custom_ops_for_converter.js'),
  }
}

...

Và đây là đoạn mã tương đương cho việc tổng hợp ( ví dụ đầy đủ ở đây ):

import alias from '@rollup/plugin-alias';

...

alias({
  entries: [
    {
      find: /@tensorflow\/tfjs$/,
      replacement: path.resolve(__dirname, './custom_tfjs/custom_tfjs.js'),
    },
    {
      find: /@tensorflow\/tfjs-core$/,
      replacement: path.resolve(__dirname, './custom_tfjs/custom_tfjs_core.js'),
    },
    {
      find: '@tensorflow/tfjs-core/dist/ops/ops_for_converter',
      replacement: path.resolve(__dirname, './custom_tfjs/custom_ops_for_converter.js'),
    },
  ],
}));

...

Nếu trình đóng gói của bạn không hỗ trợ đặt bí danh mô-đun, bạn sẽ cần thay đổi câu lệnh import của mình để nhập tensorflow.js từ custom_tfjs.js đã tạo đã được tạo ở Bước 3. Các định nghĩa Op sẽ không bị loại bỏ, nhưng hạt nhân vẫn sẽ là cây -run rẩy. Nói chung, hạt rung cây là thứ mang lại mức tiết kiệm lớn nhất về kích thước gói cuối cùng.

Nếu bạn chỉ đang sử dụng gói @tensoflow/tfjs-core thì bạn chỉ cần đặt bí danh cho gói đó.

Bước 5. Tạo gói của bạn

Chạy gói của bạn (ví dụ: webpack hoặc rollup ) để tạo gói của bạn. Kích thước của gói phải nhỏ hơn so với khi bạn chạy gói mà không có bí danh mô-đun. Bạn cũng có thể sử dụng công cụ hiển thị như thế này để xem điều gì đã đưa nó vào gói cuối cùng của bạn.

Bước 6. Kiểm tra ứng dụng của bạn

Hãy đảm bảo kiểm tra xem ứng dụng của bạn có hoạt động như mong đợi không!