How can I load/retrain/save ssd_inception_v2_coco

2019-07-08 03:31发布

问题:

ML / Tensorflow beginner.

Can any of these already-trained models be loaded on tfjs and re-trained there, then exported to Downloads or is Tensorflow python the only way to go?

I see this process is well described and documented in this tutorial for Tensorflow Python but unfortunately I can't find any documentation/tutorial to re-train an object detection model on the browser with tfjs (image classification yes, object detection no).

I see how I could load the coco-ssd model using npm, then probably even trigger saving it to downloads, but what about:

  • config file (need to modify it because I want to have only one class, not 90)
  • annotated images (both .jpg, .xml and .csv)
  • labels.pbtxt
  • .record files

Is there any way to go through the process of retraining an ssd model such as ssd_inception_v2_coco and I'm not hitting the right google keywords or is it just not possible in the current state of the framework?

回答1:

You can use transfer learning by using coco-ssd model as a feature extractor. An example of transfer-learning can be seen here.

Here is a model which extracts features using a features extractor as an input for a new sequential model.

const loadModel = async () => {
  const loadedModel = await tf.loadModel(MODEL_URL)
  console.log(loadedModel)
  // take whatever layer except last output
  loadedModel.layers.forEach(layer => console.log(layer.name))
  const layer = loadedModel.getLayer(LAYER_NAME)
  return tf.model({ inputs: loadedModel.inputs, outputs: layer.output });
}
loadModel().then(featureExtractor => {
  model = tf.sequential({
    layers: [
      // Flattens the input to a vector so we can use it in a dense layer. While
      // technically a layer, this only performs a reshape (and has no training
      // parameters).
      // slice so as not to take the batch size
      tf.layers.flatten(
        { inputShape: featureExtractor.outputs[0].shape.slice(1) }),
      // add all the layers of the model to train
      tf.layers.dense({
        units: UNITS,
        activation: 'relu',
        kernelInitializer: 'varianceScaling',
        useBias: true
      }),
      // Last Layer. The number of units of the last layer should correspond
      // to the number of classes to predict.
      tf.layers.dense({
        units: NUM_CLASSES,
        kernelInitializer: 'varianceScaling',
        useBias: false,
        activation: 'softmax'
      })
    ]
  });
})

To detect a single object out of the 90 classes of coco-ssd, one could simply use a conditional test on the prediction of coco-ssd.

const image = document.getElementById(id)

cocoSsd.load()
  .then(model => model.detect(image))
  .then(prediction => {
if (prediction.class === OBJECT_DETECTED) {
  // display it the bbox to the user}
})

If the class does not exist in coco-ssd, then one needs to builds a detector.