Tensorflow Hub和JS:如何微调预训练的模型并将其导出以在Tensorflow.js中使用?

时间:2019-01-25 09:29:36

标签: python tensorflow tensorflow.js tensorflow-hub

我试图在Tensorflow Hub上应用转移学习,并将经过训练的模型导出到Tensorflow.js。

但是,导出的训练模型只有2MB,远远不够。

似乎Tensorflow Hub模块未导出。如何导出?

def mobilenet_model_fn(features, labels, mode):
  module = hub.Module("https://tfhub.dev/google/imagenet/mobilenet_v2_100_224/classification/2")
  input_layer = features['images']  # adjust_image(features["x"])
  outputs = module(input_layer)

  logits = tf.layers.dense(inputs=outputs, units=4, name='hand_output')

  predictions = {
      # Generate predictions (for PREDICT and EVAL mode)
      "classes": tf.argmax(input=logits, axis=1),
      # Add `softmax_tensor` to the graph. It is used for PREDICT and by the
      # `logging_hook`.
      "probabilities": tf.nn.softmax(logits, name="softmax_tensor")
  }

  if mode == tf.estimator.ModeKeys.PREDICT:
      return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)

  labels_converted = tf.argmax(labels, axis=-1)

  loss = tf.losses.sparse_softmax_cross_entropy(
      labels=labels_converted, logits=logits)

  # Configure the Training Op (for TRAIN mode)
  if mode == tf.estimator.ModeKeys.TRAIN:
      optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001)
      train_op = optimizer.minimize(
          loss=loss,
          global_step=tf.train.get_global_step())
      return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)

  # Add evaluation metrics (for EVAL mode)
  accuracy = tf.metrics.accuracy(
      labels=labels_converted, predictions=predictions["classes"])
  eval_metric_ops = {"accuracy": accuracy}
  tf.summary.scalar('accuracy', accuracy[1])

  logging_hook = tf.train.LoggingTensorHook({"loss": loss,
                                            "accuracy": accuracy}, every_n_iter=2)

  return tf.estimator.EstimatorSpec(
      mode=mode, loss=loss, eval_metric_ops=eval_metric_ops, training_hooks=[logging_hook])

1 个答案:

答案 0 :(得分:0)

您正在尝试使用移动网络作为特征提取器。您有一个有关如何使用预训练的移动网络模型的示例here