[962K -> 347K] TensorflowJS 基于 Runtime 结果的 TreeShaking

基于ES6模块的 Tree Shaking，TensorFlow.js 2.x 开始就能看到package.json里sideEffect配置了，而 TensorFlow.js 3.0 新增了通过tf.profile得到Runtime 所需kernels，然后生成Custom Module实现更进一步Tree Shaking

测试使用的是face-landmarks-detection模型

变量设置

0. treeshake (true/false)
1. import full / partial
   import * as tf from 'tfjs'(full)
   import * as tf from 'tfjs-core'(partial)
   import '@tensorflow/tfjs-backend-cpu'(partial)
   import '@tensorflow/tfjs-backend-webgl'(partial)
2. custom mobule (true/false)

Type	Size
full	962kb
full + treeshake	952kb
partial	709kb
partial + treeshake	699kb
custom + full	516kb
custom + full + treeshake	505kb
custom + partial	708kb
custom + partial + treeshake	697kb

这里可以发现体积最小的是custom + full + treeshake，而不是custom + partial + treeshake，这是因为custom module已经注入对应backend所依赖的kernel

所以backend的依赖可删除

import * as tf from '@tensorflow/tfjs-core';
/  import '@tensorflow/tfjs-backend-webgl';
// import '@tensorflow/tfjs-backend-cpu';

Type	Size
custom + partial(core only)	357kb
custom + partial(core only) + treeshake	347kb

如何制作 Custom Module

Generating size-optimized browser bundles with TensorFlow.js

第一步搜集依赖

tf.profile(async () => {
  await model.estimateFaces()
}).then(e => {
  console.log(e.kernelNames);
});

编辑tfjs-config.json

kernels填入打印出的数组，models填入加载的模型的model.json

{
  "kernels": [
    "FromPixels",
    "Cast",
    "ExpandDims",
    "RotateWithOffset",
    "CropAndResize",
    "RealDiv",
    "FusedConv2D",
    "fusedConv2d__op",
    "DepthwiseConv2dNative",
    "Add",
    "Prelu",
    "MaxPool",
    "PadV2",
    "Reshape",
    "Identity",
    "Sigmoid"
  ],
  "backends": ["cpu", "webgl"],
  "models": ["./models/facemesh/model.json", "./models/blazeface/model.json"],
  "outputPath": "./custom_tfjs",
  "forwardModeOnly": true
}

配置alias

这里构建使用rollup

alias({
  entries: [
    {
      find: /@tensorflow\/tfjs$/,
      replacement: path.resolve(__dirname, options.customTfjsPath),
    },
    {
      find: /@tensorflow\/tfjs-core$/,
      replacement: path.resolve(__dirname, options.customTfjsCorePath),
    },
    {
      find: '@tensorflow/tfjs-core/dist/ops/ops_for_converter',
      replacement: path.resolve(__dirname, options.customOpsPath),
  ],
}),

生成Custom Module

package.json

{
  "make-custom": "node ./node_modules/@tensorflow/tfjs/dist/tools/custom_module/cli.js --config tfjs_config.json",
}

运行测试

这里只有custom-full和custom-full-treeshake不能正常工作，而tfjs_custom_module_demo里blazeface是可以正常工作，所以建议custom + partial(core only) + treeshake

测试demo Github

参考

0. Generating size-optimized browser bundles with TensorFlow.js
1. tfjs_custom_module_demo