tensorflow.js: Unable to load a model saved in tfjs

Question

I am unable to load a sequential model I created using @tensorflow/tfjs-node back into node. My model creation looks like this:

import * as tf from '@tensorflow/tfjs-node';
      
model = tf.sequential();

// Add a single input layer
model.add(
   tf.layers.dense({ inputShape: [1.0], units: 1, useBias: true })
);

model.add(
   tf.layers.dense({
      inputShape: [1],
      units: 10,
      activation: 'relu',
      kernelRegularizer: tf.regularizers.l2({ l2: 0.01 }),
      useBias: true
   })
);

// Add additional hidden layers with L2 regularization
model.add(
   tf.layers.dense({
      units: 8,
      activation: 'relu',
      kernelRegularizer: tf.regularizers.l2({ l2: 0.01 }),
      useBias: true
   })
);

model.add(
   tf.layers.dense({
      units: 6,
      activation: 'relu',
      kernelRegularizer: tf.regularizers.l2({ l2: 0.01 }),
      useBias: true
   })
);

// Add an output layer with linear activation
model.add(
   tf.layers.dense({ units: 1, activation: 'linear', useBias: true })
);

// Add an output layer
model.add(tf.layers.dense({ units: 1, useBias: true }));

I am saving the model like this (bit crude at the moment, but it seems otherwise to work):

const createSavableModel = async () => {
   await model.save(tf.io.withSaveHandler(artifacts => {
      modelArtifacts = serialize.serialize(artifacts);
      return Promise.resolve({
         modelArtifactsInfo: {
            dateSaved: new Date()
         } as any
      });
   }));
};

and them I'm stashing away the modelArtifacts in MongoDB.

Then, when trying to load the model, like this:

model = await tf.loadLayersModel(tf.io.fromMemory(existingModelArtifacts));

I get a very long error:

Error: Unknown layer: {"modelTopology":{"className"...

1. The layer is defined in Python, in which case it needs to be ported to TensorFlow.js or your JavaScript code.
2. The custom layer is defined in JavaScript, but is not registered properly with tf.serialization.registerClass().

Now the model was created entirely in tensorflow.js, and the error doesn't really tell me which layer it's unhappy with.

What might be causing this?

EDIT

I have reduced my model layers:

model = tf.sequential();

// Add a single input layer
model.add(
   tf.layers.dense({ inputShape: [1.0], units: 1, useBias: true })
);

model.add(tf.layers.dense({ units: 1, useBias: true }));

and here is the topology:

{
  "class_name": "Sequential",
  "config": {
    "name": "sequential_1",
    "layers": [
      {
        "class_name": "Dense",
        "config": {
          "units": 1,
          "activation": "linear",
          "use_bias": true,
          "kernel_initializer": {
            "class_name": "VarianceScaling",
            "config": {
              "scale": 1,
              "mode": "fan_avg",
              "distribution": "normal",
              "seed": null
            }
          },
          "bias_initializer": {
            "class_name": "Zeros",
            "config": {}
          },
          "kernel_regularizer": null,
          "bias_regularizer": null,
          "activity_regularizer": null,
          "kernel_constraint": null,
          "bias_constraint": null,
          "name": "dense_Dense1",
          "trainable": true,
          "batch_input_shape": [
            null,
            1
          ],
          "dtype": "float32"
        }
      },
      {
        "class_name": "Dense",
        "config": {
          "units": 1,
          "activation": "linear",
          "use_bias": true,
          "kernel_initializer": {
            "class_name": "VarianceScaling",
            "config": {
              "scale": 1,
              "mode": "fan_avg",
              "distribution": "normal",
              "seed": null
            }
          },
          "bias_initializer": {
            "class_name": "Zeros",
            "config": {}
          },
          "kernel_regularizer": null,
          "bias_regularizer": null,
          "activity_regularizer": null,
          "kernel_constraint": null,
          "bias_constraint": null,
          "name": "dense_Dense2",
          "trainable": true
        }
      }
    ]
  },
  "keras_version": "tfjs-layers 4.2.0",
  "backend": "tensor_flow.js"
}

So why can't I load this with tf.loadLayersModel(tf.io.fromMemory(existingModelArtifacts))?

Answer 1

Ok, so I decided to go via the filesystem:

const createArchive = async (sourceDir: string): Promise<Buffer> => {
   const archive = archiver('zip', { zlib: { level: 9 } });
   const chunks: Buffer[] = [];

   return new Promise<Buffer>((resolve, reject) => {
      archive.on('data', (data) => {
         chunks.push(data);
      });

      archive.on('end', () => {
         resolve(Buffer.concat(chunks));
      });

      archive.directory(sourceDir, false);
      archive.finalize();
   });
};

const extractBuffer = (buffer: Buffer, destinationDir: string): Promise<void> => {
   return new Promise<void>((resolve, reject) => {
      const extractor = unzipper.Extract({ path: destinationDir });

      extractor.on('error', (err) => {
         reject(`Error extracting archive: ${err.message}`);
      });

      extractor.on('close', () => {
         console.log(`Archive extracted to: ${destinationDir}`);
         resolve();
      });

      extractor.write(buffer);
      extractor.end();
   });
};

const modelToBuffer = async (model: any) => {
   const modelSaveTempDir = await mkdtemp(path.join(os.tmpdir(), `tf-${pairName}`));
   await model.save(`file://${modelSaveTempDir}`);

   return createArchive(modelSaveTempDir);
};

const bufferToModel = async (buffer: Buffer): Promise<any> => {
   const modelTempDir = await mkdtemp(path.join(os.tmpdir(), `tf-${pairName}`));
   await extractBuffer(buffer, modelTempDir);
   return await tf.loadLayersModel(`file://${modelTempDir}/model.json`);
};