使用 TensorFlow.js 生成尺寸优化的浏览器包

概述

TensorFlow.js 3.0 支持构建尺寸优化、面向生产的浏览器包。换句话说,我们希望让您更轻松地将更少的 JavaScript 代码发送到浏览器。

此功能面向具有生产用例的用户,这些用户将特别受益于减少其有效负载的字节数(因此愿意付出努力来实现这一点)。要使用此功能,您应该熟悉ES 模块、JavaScript 打包工具(如webpackrollup)以及摇树优化/死代码消除等概念。

本教程演示了如何创建一个自定义 tensorflow.js 模块,该模块可与打包器一起使用,以生成使用 tensorflow.js 的程序的尺寸优化构建。

术语

在本文件中,我们将使用一些关键术语

ES 模块 - 标准 JavaScript 模块系统。在 ES6/ES2015 中引入。可以通过使用importexport 语句来识别。

打包 - 将一组 JavaScript 资产分组/打包成一个或多个可在浏览器中使用的 JavaScript 资产。这是通常生成最终资产的步骤,这些资产将提供给浏览器。应用程序通常会直接从转译的库源代码中进行自己的打包常见的打包器包括rollupwebpack。打包的最终结果称为(或有时称为,如果它被拆分为多个部分)。

摇树优化/死代码消除 - 删除最终编写的应用程序未使用的代码。这在打包过程中完成,通常在缩小步骤中完成。

操作 (Ops) - 对一个或多个张量执行的数学运算,该运算生成一个或多个张量作为输出。操作是“高级”代码,可以使用其他操作来定义其逻辑。

内核 - 与特定硬件功能绑定的操作的特定实现。内核是“低级”且特定于后端的。一些操作从操作到内核具有单对一映射,而其他操作使用多个内核。

范围和用例

仅推理的图模型

我们从用户那里听到的关于此的主要用例,以及我们在本次发布中支持的用例是使用TensorFlow.js 图模型进行推理。如果您使用的是TensorFlow.js 层模型,则可以使用tfjs-converter 将其转换为图模型格式。图模型格式对于推理用例更有效。

使用 tfjs-core 进行低级张量操作

我们支持的另一个用例是直接使用 @tensorflow/tjfs-core 包进行低级张量操作的程序。

我们对自定义构建的方法

我们在设计此功能时遵循的核心原则包括以下内容

  • 最大限度地利用 JavaScript 模块系统 (ESM),并允许 TensorFlow.js 的用户也这样做。
  • 使 TensorFlow.js 尽可能地通过现有打包器(例如 webpack、rollup 等)进行摇树优化。这使用户能够利用这些打包器的所有功能,包括代码拆分等功能。
  • 尽可能地为对包大小不那么敏感的用户保持易用性。这确实意味着生产构建将需要更多努力,因为我们库中的许多默认设置都支持易用性而不是尺寸优化构建。

我们工作流程的主要目标是为 TensorFlow.js 生成一个自定义JavaScript 模块,该模块仅包含我们尝试优化的程序所需的函数。我们依靠现有的打包器来完成实际的优化。

虽然我们主要依靠 JavaScript 模块系统,但我们还提供了一个自定义CLI 工具来处理用户界面代码中难以通过模块系统指定的某些部分。以下是两个示例

  • 存储在model.json 文件中的模型规范
  • 我们使用的操作到后端特定内核的调度系统。

这使得生成自定义 tfjs 构建比仅仅将打包器指向常规 @tensorflow/tfjs 包更复杂。

如何创建尺寸优化的自定义包

步骤 1:确定您的程序正在使用哪些内核

此步骤允许我们确定您在给定所选后端的情况下运行的任何模型或预/后处理代码使用的所有内核。

使用 tf.profile 运行使用 tensorflow.js 的应用程序部分并获取内核。它看起来像这样

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

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

将该内核列表复制到剪贴板,以便下一步使用。

您需要使用与您要在自定义包中使用的相同后端对代码进行分析。

如果您的模型发生变化或您的预/后处理代码发生变化,您将需要重复此步骤。

步骤 2. 为自定义 tfjs 模块编写配置文件

这是一个示例配置文件。

它看起来像这样

{
  "kernels": ["Reshape", "_FusedMatMul", "Identity"],
  "backends": [
      "cpu"
  ],
  "models": [
      "./model/model.json"
  ],
  "outputPath": "./custom_tfjs",
  "forwardModeOnly": true
}
  • kernels: 要包含在包中的内核列表。从步骤 1 的输出中复制此内容。
  • backends: 您要包含的后端列表。有效选项包括“cpu”、“webgl”和“wasm”。
  • models: 应用程序中加载的模型的 model.json 文件列表。如果您的程序不使用 tfjs_converter 加载图形模型,则可以为空。
  • outputPath: 用于放置生成模块的文件夹的路径。
  • forwardModeOnly: 如果您想为之前列出的内核包含梯度,请将其设置为 false。

步骤 3. 生成自定义 tfjs 模块

使用配置文件作为参数运行自定义构建工具。您需要安装 @tensorflow/tfjs 包才能访问此工具。

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

这将在 outputPath 中创建一个包含一些新文件的文件夹。

步骤 4. 配置您的打包器将 tfjs 链接到新的自定义模块。

在 webpack 和 rollup 等打包器中,我们可以将对 tfjs 模块的现有引用链接到我们新生成的自定义 tfjs 模块。为了最大限度地减少包大小,需要链接三个模块。

以下是 webpack 中的代码片段(完整示例在此处

...

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'),
  }
}

...

以下是 rollup 的等效代码片段(完整示例在此处

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'),
    },
  ],
}));

...

如果您的打包器不支持模块链接,您需要更改您的 import 语句,以便从步骤 3 中创建的生成的 custom_tfjs.js 中导入 tensorflow.js。操作定义不会被树形抖动,但内核仍然会被树形抖动。通常,树形抖动内核是最终包大小节省最大的地方。

如果您只使用 @tensoflow/tfjs-core 包,那么您只需要链接该包。

步骤 5. 创建您的包

运行您的打包器(例如 webpackrollup)来生成您的包。包的大小应该小于在没有模块链接的情况下运行打包器时的包大小。您还可以使用 等可视化工具来查看最终包中包含的内容。

步骤 6. 测试您的应用程序

确保您的应用程序按预期工作!