javascriptnode.jswebpacktensorflow.js

Load appropriate TensorFlow.js version for app or Node script


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.


Solution

  • 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"
       }