Support aggressive tree-shaking to remove code not needed to execute a model for inference.

[PiMastah]

Describe the problem or feature request

Let me start out by appreciating all the effort that is going into this project!

I am currently trying to build an agent in order to play a game on a coding site. Their limit for the size of your uploaded code seems to be around 100 kb. Using webpack and including the whole library, I am already at around 1,7 mb. I was therefore wondering if it is possible to export a standalone activation function of a given model in order to save some bytes by not including the whole library? Looking through the documentation, I did not find a way to do just this but I remember that e.g. synaptic allows to do exactly that.

Thanks in advance for your effort!

[bileschi]

Hi PiMastah,

It sounds like you are looking for a way to export a library that contains just enough code & data to run my_model.predict to minimize the footprint of the code. Is that right? Something like a TensorFlow Lite for TensorFlow.js?

[PiMastah]

Yes, @bileschi - what you describe would fit my needs very well.

[bileschi]

Hrm., it looks like tf.Lite is 300KB minimum, so that probably won't meet your needs. I expect at 100kb you would be looking at a pretty small model. Do you have a model which meets your performance quality expectations, or are you more at the technology exploratory phase?

[PiMastah]

It is definitely the latter at the moment. If it can't be done, nothing much is lost. I know the limitations are pretty severe (after talking to one of the platform devs, the upload limit is actually 100k characters) but I was more trying to picture the whole task as a maybe somewhat unusual challenge which would however allow me to dive deep into a lot of things related to TF and thus serve as a learning opportunity :) I will still be using TF.js to build a sophisticated model which will run locally either way.

[nsthorat]

This is something we've thought about and might pursue in the future, so keeping this open. For now we don't have any near term plans to do this.

[mohammedzamakhan]

@igorminar @stephenfluin can probably guide you guys. Angular Team has done a great job, and tree shaking and ahead of time compilation has helped a lot of Angular Projects

[fredyagomez]

+1. I think this is really important if anybody wants to consume this. No as a final consumer but as a vendor or middle consumer.

import { tensor2d } from "@tensorflow/tfjs"; imports the whole library! 2.6MB. @nsthorat why the low priority for this?

ok, after looking at the dependencies I ended up doing this one. import { tensor2d } from '@tensorflow/tfjs-core'; which is 1.37 MB.

[ydennisy]

Hello All,

Any idea why tree shaking as implemented by webpack does not work for tfjs-core?

Do you guys plan to support the mentioned use case of picking and choosing objects from the library?

import { tensor2d } from '@tensorflow/tfjs-core';

Thanks in advance!

[ydennisy]

Any update here?

[nsthorat]

We’re starting to work on this this quarter!

[bwnodak]

How's this coming along? Would love to start using this at scale soon. Thanks so much for the effort!

[tafsiri]

We are working on it! I hesitate to give a timeline right now as its quite a big change that involves modularizing all our ops and kernels as well as some internal architecture changes. This github label will track issues/PRs related to this.

[ydennisy]

Hi @tafsiri @nsthorat I can see work on some of the issues labelled - would you happen to have an estimate yet?

Or is there any pre-release stuff one can play with?

[tafsiri]

@ydennisy we hope to have something out by the end of the year, at the very least in pre-release form if not finalised.

As for pre-release stuff, we don't have any documentation yet, but if you are adventurous and don't mind looking at source code/tests, you can see the workflow we are targeting here https://github.com/tensorflow/tfjs/tree/master/e2e/custom_bundle/dense_model. The one tip I can give if you look at that is that we use the result of tf.profile to get the list of kernels in the config file.

[tafsiri]

To anyone interested in this topic and adventurous enough to try out a pre-release and give us feedback! We have a release candidate for a version of tfjs with improved support and workflows for tree shaking. Check out the release notes here which has more of the details on the new functionality.

Thanks!

[LukeJVinton]

Hi @tafsiri. Thanks for this pre-release, it looks really promising!

I am attempting to create a size optimised bundle following the guide. However, I am getting errors when generating the custom tfjs module in step 3. This is my terminal output:

`$ npx tfjs-custom-bundle --config custom_tfjs_config.json npm ERR! code E404 npm ERR! 404 Not Found - GET https://registry.npmjs.org/tfjs-custom-bundle - Not found npm ERR! 404 npm ERR! 404 'tfjs-custom-bundle@latest' is not in the npm registry. npm ERR! 404 You should bug the author to publish it (or use the name yourself!) npm ERR! 404 npm ERR! 404 Note that you can also install from a npm ERR! 404 tarball, folder, http url, or git url.

npm ERR! A complete log of this run can be found in: npm ERR! /Users/lukevinton/.npm/_logs/2021-01-26T14_00_34_890Z-debug.log`

I have tried to search for similar problems but haven't found anything as of yet. Any help with this would be greatly appreciated, I am quite new to JS so apologies if the error is my end.

[tafsiri]

@LukeJVinton sorry about that I think the doc was a bit out of date, between rc's we renamed tfjs-custom-bundle to tfjs-custom-module. Could you try that and let me know if it works.

[tafsiri]

@LukeJVinton also do you have @tensorflow/tfjs installed?

[LukeJVinton]

@tafsiri no problem. I figured the mismatch out and was going to post to point it out. I've now managed to get past step 3, thanks for the pointers.

[tafsiri]

@LukeJVinton Thanks for reporting! Would also be curious to hear how it goes and what kind of size reduction you end up seeing.

[LukeJVinton]

@tafsiri I've hit another problem, probably something I am doing wrong.

Once I've built the custom modules and updated my webpack.config.js I run npm run build which produces a warning regarding the loadLayersModel. In the printout I get the message:

$ npm run build

> ds-inbrowser-ml@1.0.0 build
> webpack

asset main.js 77.2 KiB [emitted] [minimized] (name: main) 1 related asset
orphan modules 1.09 MiB [orphan] 388 modules
runtime modules 495 bytes 2 modules
cacheable modules 379 KiB
  ./src/index.js + 88 modules 379 KiB [built] [code generated]
  node-fetch (ignored) 15 bytes [built] [code generated]
  util (ignored) 15 bytes [built] [code generated]

WARNING in ./src/index.js 49:28-46
export 'loadLayersModel' (imported as 'tf') was not found in '@tensorflow/tfjs' (possible exports: Abs, Acos, Acosh, AdadeltaOptimizer, AdagradOptimizer, AdamOptimizer, AdamaxOptimizer, Add, AddN, All, Any, ArgMax, ArgMin, Asin, Asinh, Atan, Atan2, Atanh, AvgPool, AvgPool3D, AvgPool3DGrad, AvgPoolGrad, BatchMatMul, BatchToSpaceND, Bincount, BroadcastTo, Cast, Ceil, ClipByValue, Complex, ComplexAbs, Concat, Conv2D, Conv2DBackpropFilter, Conv2DBackpropInput, Conv3D, Conv3DBackpropFilterV2, Conv3DBackpropInputV2, Cos, Cosh, CropAndResize, Cumsum, DataStorage, DenseBincount, DepthToSpace, DepthwiseConv2dNative, DepthwiseConv2dNativeBackpropFilter, DepthwiseConv2dNativeBackpropInput, Diag, Dilation2D, Dilation2DBackpropFilter, Dilation2DBackpropInput, ENV, Elu, EluGrad, Environment, Equal, Erf, Exp, ExpandDims, Expm1, FFT, Fill, FlipLeftRight, Floor, FloorDiv, FromPixels, FusedBatchNorm, FusedConv2D, FusedDepthwiseConv2D, GatherNd, GatherV2, Greater, GreaterEqual, IFFT, Identity, Imag, IsFinite, IsInf, IsNan, KernelBackend, LRN, LRNGrad, LeakyRelu, Less, LessEqual, LinSpace, Log, Log1p, LogSoftmax, LogicalAnd, LogicalNot, LogicalOr, MathBackendCPU, Max, MaxPool, MaxPool3D, MaxPool3DGrad, MaxPoolGrad, MaxPoolWithArgmax, Maximum, Mean, Min, Minimum, MirrorPad, Mod, MomentumOptimizer, Multinomial, Multiply, Neg, NonMaxSuppressionV3, NonMaxSuppressionV4, NonMaxSuppressionV5, NotEqual, OP_SCOPE_SUFFIX, OneHot, OnesLike, Optimizer, Pack, PadV2, Pool, Pow, Prelu, Prod, RMSPropOptimizer, Range, Rank, Real, RealDiv, Reciprocal, Reduction, Relu, Relu6, Reshape, ResizeBilinear, ResizeBilinearGrad, ResizeNearestNeighbor, ResizeNearestNeighborGrad, Reverse, RotateWithOffset, Round, Rsqrt, SGDOptimizer, ScatterNd, Select, Selu, Sigmoid, Sign, Sin, Sinh, Slice, Softmax, Softplus, SpaceToBatchND, SparseToDense, SplitV, Sqrt, Square, SquaredDifference, Step, StridedSlice, Sub, Sum, Tan, Tanh, Tensor, TensorBuffer, Tile, TopK, Transpose, Unique, Unpack, UnsortedSegmentSum, Variable, ZerosLike, _FusedMatMul, abs, acos, acosh, add, addN, all, any, argMax, argMin, asin, asinh, atan, atan2, atanh, avgPool, avgPool3d, backend, backend_util, basicLSTMCell, batchNorm, batchNorm2d, batchNorm3d, batchNorm4d, batchToSpaceND, bincount, booleanMaskAsync, broadcastTo, browser, buffer, cast, ceil, clipByValue, clone, complex, concat, concat1d, concat2d, concat3d, concat4d, conv1d, conv2d, conv2dTranspose, conv3d, conv3dTranspose, copyRegisteredKernels, cos, cosh, cosineWindow, cumsum, customGrad, denseBincount, deprecationWarn, depthToSpace, depthwiseConv2d, device_util, diag, dilation2d, disableDeprecationWarnings, dispose, disposeVariables, div, divNoNan, dot, dropout, elu, enableDebugMode, enableProdMode, enclosingPowerOfTwo, engine, env, equal, erf, exp, expandDims, expm1, eye, fft, fill, findBackend, findBackendFactory, floor, floorDiv, fused, gather, gatherND, gather_util, getBackend, getGradient, getKernel, getKernelsForBackend, grad, grads, greater, greaterEqual, ifft, imag, image, inTopKAsync, io, irfft, isFinite, isInf, isNaN, keep, kernel_impls, leakyRelu, less, lessEqual, linalg, linspace, localResponseNormalization, log, log1p, logSigmoid, logSoftmax, logSumExp, logicalAnd, logicalNot, logicalOr, logicalXor, losses, matMul, math, max, maxPool, maxPool3d, maxPoolWithArgmax, maximum, mean, memory, min, minimum, mirrorPad, mod, moments, movingAverage, mul, multiRNNCell, multinomial, neg, nextFrame, norm, notEqual, oneHot, ones, onesLike, op, outerProduct, pad, pad1d, pad2d, pad3d, pad4d, pool, pow, prelu, print, prod, profile, rand, randomGamma, randomNormal, randomUniform, range, ready, real, reciprocal, registerBackend, registerGradient, registerKernel, relu, relu6, removeBackend, reshape, reverse, reverse1d, reverse2d, reverse3d, reverse4d, rfft, round, rsqrt, scalar, scatterND, scatter_util, selu, separableConv2d, serialization, setBackend, setPlatform, setdiff1dAsync, shared, sigmoid, sign, signal, sin, sinh, slice, slice1d, slice2d, slice3d, slice4d, slice_util, softmax, softplus, spaceToBatchND, sparseToDense, spectral, split, sqrt, square, squaredDifference, squeeze, stack, step, stridedSlice, sub, sum, sumOutType, tan, tanh, tensor, tensor1d, tensor2d, tensor3d, tensor4d, tensor5d, tensor6d, tensor_util, test_util, tidy, tile, time, topk, train, transpose, truncatedNormal, unique, unregisterGradient, unregisterKernel, unsortedSegmentSum, unstack, upcastType, util, valueAndGrad, valueAndGrads, variable, variableGrads, version_core, version_cpu, where, whereAsync, zeros, zerosLike)

I am wondering why the loadLayersModel didn't get picked up by tf.profile. In my index.js file I have moved all the tensorflow dependant code into one function so that it is easy to profile in one go, my function looks like this:

    const xTensor = tf.tensor2d([x])
    const yTensor = tf.tensor2d([y])
    const xyTensor = tf.stack([xTensor, yTensor], 2)

    const seq_model = await tf.loadLayersModel(get_model_path())

    // save model to local storage if not already there
    if (get_model_path() != 'localstorage://model') {
        await seq_model.save('localstorage://model')
    };

    var predictedBehaviourTensor = seq_model.predict(xyTensor)
    var predNum = predictedBehaviourTensor.dataSync()[0]

    return predNum
};

and I apply tf.profile like this:

const profileInfo = await tf.profile(async () => {
    // You must profile all uses of tf symbols.
    predNum = await predict_behaviour(tempX, tempY)
});

At the top of the file I have the depency imports:

import * as tf from '@tensorflow/tfjs';
import * as tfc from '@tensorflow/tfjs-core';

Apologies for the essay, please let me know if there is info I have overlooked that would be useful.

[tafsiri]

To use the custom module flow you need to convert your layers model into a graph model. You can do that with tfjs-converter

[LukeJVinton]

My model contains an LSTM layer which the conversion to graph model does not support yet (https://github.com/tensorflow/tfjs/tree/master/tfjs-converter#python-to-javascript). That's the only reason I use a layers model. Is there any way around this other than creating a model without recurrent layers?

[tafsiri]

You could try manually importing tfjs-layers e.g. import * as tfl from @tensorflow/tfjs-layers. That will give you access to loadLayers method , however since this flow isn't the intended path, I can't comment on how much bundle size saving you can expect (it wont be as good as inference focused graphmodels) layers also includes a number of things for training. You can try this out but is pretty much untested.

[LukeJVinton]

@tafsiri Thanks for all your help. I have managed to create a custom module using the guide. For now I have opted to use a model that does not include recurrent layers as the bundle size is critical. At some point it would be great to be able to follow this process using a model containing recurrent layers.

The total bundle size I am getting for the tfjs dependencies is now about 5kB! The breakdown is:

-rw-r--r--  1 lukevinton  staff   1.0K 28 Jan 11:46 custom_ops_for_converter.js
-rw-r--r--  1 lukevinton  staff   3.1K 28 Jan 11:46 custom_tfjs.js
-rw-r--r--  1 lukevinton  staff   904B 28 Jan 11:46 custom_tfjs_core.js

By the way, I noticed that the links in step 4 point towards pages that no longer exist. I think the urls should be edited by replacing "bundle" with "module" as such:

• Change https://github.com/tensorflow/tfjs/blob/master/e2e/custom_bundle/dense_model/webpack.config.js to https://github.com/tensorflow/tfjs/blob/master/e2e/custom_module/dense_model/webpack.config.js

• Change https://github.com/tensorflow/tfjs/blob/master/e2e/custom_bundle/dense_model/rollup.config.js to https://github.com/tensorflow/tfjs/blob/master/e2e/custom_module/dense_model/rollup.config.js

[tafsiri]

@LukeJVinton thats awesome that you got it working. However I should say that those files just import the parts of tfjs that will end up in your bundle. You will need to use a bundle analyzer to measure the final size of what gets pulled in by webpack after tree-shaking. Thanks for reporting the broken links, will update that shortly.

[LukeJVinton]

@tafsiri Hi there.

I've found that this size optimisation process has worked when the model is saved to the browsers "localStorage". However I run into problems when trying to grab a model from an AWS bucket whilst using the size optimised tfjs packages.

Without size optimisation I am able to fetch the model from an AWS bucket. Once I implement size optimisation (i.e go over the steps from the size optimisation guide) I see this error in the console when attempting to load the model:

Uncaught (in promise) TypeError: Cannot read property 'fetch' of undefined
  | HTTPRequest | @ | http.js:54
-- | -- | -- | --
  | http | @ | http.js:337
  | httpRouter | @ | http.js:259
  | eval | @ | router_registry.js:80
  | getHandlers | @ | router_registry.js:79
  | getLoadHandlers | @ | router_registry.js:72
  | getLoadHandlers | @ | router_registry.js:91
  | findIOHandler | @ | graph_model.js:96
  | load | @ | graph_model.js:114
  | loadGraphModel | @ | graph_model.js:394
  | predict_behaviour | @ | index.js:64
  | async function (async) |   |  
  | eval | @ | index.js:114

Please let me now if I can provide more information or should be posting in a more suitable forum.

----- Update -----

After a bit of debugging its seems that the error occurs at at this point: https://github.com/tensorflow/tfjs/blob/1edd7170aae8da3d3ab15a2dcd172c790d8d98bd/tfjs-core/src/io/http.ts#L64

To test I've added the following logging lines to my index.js file:

    console.log("navigator product: ", navigator.product)
    console.log(tfcore.env())
    console.log(tfcore.env().platform)

Without using the size optimised custom module I get these printouts in the console:

navigator product:  Gecko
index.js:57 Environment {global: Window, flags: {…}, flagRegistry: {…}, urlFlags: {…}, platformName: "browser", …}
index.js:58 PlatformBrowser {}

Whereas, using the custom tfjs module I get these printouts in the console:

navigator product:  Gecko
index.js:59 Environment {global: Window, flags: {…}, flagRegistry: {…}, urlFlags: {…}}
index.js:60 undefined

It looks like the platform properties of tfcore.env() are not being set when switching to use the size optimised custom module. Adding this line to my index.js file seems to provide a workaround: tf.env().setPlatform('browser', new PlatformBrowser());

[google-ml-butler[bot]]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed in 7 days if no further activity occurs. Thank you.

[google-ml-butler[bot]]

Closing as stale. Please @mention us if this needs more attention.