javascripttensorflow.jstfjs-node

How to .apply() a layer to the outputs of a model (transfer learning)


I am trying to finetune a CNN in tensorflow.js. To do this, I would like to add a head to the final layer of the pretrained model. The equivalent code in python tensorflow is as follows, where we add an average pooling layer to the pretrained efficientnet.

import tensorflow as tf

img_base = tf.keras.applications.efficientnet.EfficientNetB0(include_top=False, weights='imagenet')
img_base = tf.keras.layers.GlobalAveragePooling2D()(img_base.output)

However, the same code in JavaScript results in an error.

const tf = require('@tensorflow/tfjs-node');

const getModel = async function () {
    const imgBase = await tf.loadLayersModel('file://./tfjs_models/efficientnetb0_applications_notop/model.json');
    const imgPoolLayer = tf.layers.globalAveragePooling2d({dataFormat: 'channelsLast'});
    const imgPool = imgPoolLayer.apply(imgBase.outputs);
}
Error: Arguments to apply() must be all SymbolicTensors or all Tensors
    at new ValueError (/home/stanleyzheng/kds/kds-melanoma/tfjs_scripts/node_modules/@tensorflow/tfjs-layers/dist/tf-layers.node.js:16792:28)
    at Concatenate.Layer.apply (/home/stanleyzheng/kds/kds-melanoma/tfjs_scripts/node_modules/@tensorflow/tfjs-layers/dist/tf-layers.node.js:19983:19)
    at getModel (/home/stanleyzheng/kds/kds-melanoma/tfjs_scripts/train.js:32:29)

Printing imgBase.outputs gives us the following result. imgBase.outputs[0] returns the same error as above.

[
  SymbolicTensor {
    dtype: 'float32',
    shape: [ null, 1280, 7, 7 ],
    sourceLayer: Activation {
      _callHook: null,
      _addedWeightNames: [],
      _stateful: false,
      id: 236,
      activityRegularizer: null,
      inputSpec: null,
      supportsMasking: true,
      _trainableWeights: [],
      _nonTrainableWeights: [],
      _losses: [],
      _updates: [],
      _built: true,
      inboundNodes: [Array],
      outboundNodes: [],
      name: 'top_activation',
      trainable_: true,
      initialWeights: null,
      _refCount: 1,
      fastWeightInitDuringBuild: true,
      activation: Swish {}
    },
    inputs: [ [SymbolicTensor] ],
    callArgs: {},
    outputTensorIndex: undefined,
    id: 548,
    originalName: 'top_activation/top_activation',
    name: 'top_activation/top_activation',
    rank: 4,
    nodeIndex: 0,
    tensorIndex: 0
  }
]

How can we get the outputs of a base model such that it can be input into a separate layer? Thanks.


Solution

  • Well, turns out, with a minimal example, it works. Simply inputting model.outputs into layer.apply() works.

    const tf = require(@tensorflow/tfjs)
    
    const getModel = async function () {
        const baseModel = await tf.sequential();
        baseModel.add(tf.layers.conv2d({inputShape: [28, 28, 1], kernelSize: 5, filters: 8, strides: 1, activation: 'relu'}))
        let imgPoolLayer = tf.layers.globalAveragePooling2d({dataFormat: 'channelsLast'});
        let imgPool = imgPoolLayer.apply(baseModel.outputs);
    }
    getModel()