I have a TensorFlow.js model that I want to use both in the browser and from a script run at the command line via Node. The model is defined in a CommonJS module, so I can easily load it in both environments. I am using Webpack to bundle the browser app. But I can't figure out how to load the appropriate version of TensorFlow.js. I want to use tfjs-node
when running the script, but tfjs
in the browser.
I tried using dynamic import to load the right module at runtime, but this forces me to use a CDN for tfjs
, since it'll be loading it in the browser. What I want to do is build my app using tfjs
, but require
the model into an independent Node script, and use tfjs-node
in that context. This must be a reasonably common scenario, but I cannot see a way to do it. What am I missing?
Edit: I'm thinking maybe I can achieve this by bundling for both the browser and node? That's possible using the target
property in webpack.config.js
. It's possible to specify multiple targets. But how to set the version of TensorFlow.js for each target? Possibly by aliasing the relevant package within each target config? So I would just use the same alias, but it would map to tfjs
for the browser target config, and tfjs-node
in the Node target config.
OK, I eventually found a satisfactory way to achieve this using environment variables.
It's possible to pass environment variables to Webpack and access them during the compilation process, so their values can be substituted into your code. We need to modify webpack.config.js
to do this, in the following way:
const path = require('path');
const webpack = require('webpack');
module.exports = (env) => {
return {
entry: './index.js',
output: {
filename: 'bundle.js',
path: path.resolve(__dirname, 'public/dist'),
publicPath: '/dist/'
},
plugins: [
new webpack.DefinePlugin({ "process.env.TFJS": JSON.stringify(`${env.TFJS}`) })
]
}
};
We need to define the exported config object as a function, so any environment variables passed via the command line are accessible via the env
argument. Then we use Webpack's DefinePlugin
to make the environment variables available in our code at compile time. Here we are exposing the variable TFJS
, passed via the command line in package.json
like this:
"scripts": {
"build": "webpack --config webpack.config.js --env TFJS=tfjs"
}
So in the file where I define my model I require
TensorFlow like this:
const tf = require(`@tensorflow/${process.env.TFJS}`);
For the build file Webpack will substitute @tensorflow/tfjs
. The neat thing is if we require
this module directly in a node script we can access a different value for process.env.TFJS
, passed via the command line. This is what I am doing:
TFJS=tfjs-node node my-script.js
So inside the script process.env.TFJS
will be set to tfjs-node
, which will be interpolated into the require
call. It works! Neat.
The scripts
section of my package.json
now looks something like this:
"scripts": {
"build": "webpack --config webpack.config.js --env TFJS=tfjs",
"my-script": "TFJS=tfjs-node node my-script.js"
}