如何还原使用Dataset API训练的Tensorflow模型?

时间:2019-05-21 23:04:06

标签: python tensorflow restore tensorflow-datasets

我正在使用带有可迭代的迭代器的Dataset API训练我的模型,就像在Importing Data Tutorial here中一样。 问题是还原模型时。它还可以从训练中恢复其形状的手柄占位符。这意味着它希望为其提供示例和标签。

    def loadTFRecord(filenames):          
      dataset = tf.data.TFRecordDataset([filenames])
      dataset = dataset.map(extract_img_func)
      dataset = dataset.batch(batchsize)
      handle = tf.placeholder(tf.string, shape=[])
      iterator = tf.data.Iterator.from_string_handle(handle, dataset.output_types, dataset.output_shapes)
      training_iterator = dataset.make_one_shot_iterator()
      next_element = iterator.get_next() 
      training_handle = self.sess.run(training_iterator.string_handle())
      return next_element #next_element[0] is the example img, next_element[1] is the label

    def model_fn(images, labels=None, train=False):
      input_layer = images
      ...
      predictions = last_layer
      if train:
        return predictions

      # Calculate loss
      loss = tf.losses.mean_squared_error(labels, predictions)
      learning_rate = tf.train.exponential_decay(learning_rate=learningRate, staircase=True)
      optimizer = tf.train.GradientDescentOptimizer(learning_rate=learning_rate)
      train_op = optimizer.minimize(
          loss=loss,
          global_step=global_step)

      return train_op, predictions, loss

为此,我正在创建训练模型:

examples, labels = loadTFRecord("path/to/tfrecord")
model_fn(examples, labels=labels)
saver = tf.train.Saver(max_to_keep=4, keep_checkpoint_every_n_hours=0.5)
... #training here
saver.save(sess, "path/to/")

现在的问题是,当我想还原模型以进行推理时。 我想做的是还原模型并传递另一个可加载的迭代器,该迭代器从磁盘加载一些.png文件。我这样做类似于加载TFRecord文件。

def load_images(filenames):
  dataset = tf.data.Dataset.from_tensor_slices(filenames)
  dataset = dataset.map(lambda x: tf.image.resize_images(self.normalize(tf.image.decode_png(tf.read_file(x), channels = 3)), [IM_WIDTH, IM_HEIGHT]))
  dataset = dataset.batch(1)
  iterator = tf.data.Iterator.from_string_handle(handle, dataset.output_types, dataset.output_shapes)
  iterator = dataset.make_one_shot_iterator()
  next_img = iterator.get_next()
  training_handle = sess.run(iterator.string_handle())
  return next_img

现在的问题是,将其传递给还原的模型时是这样的:

  saver = tf.train.import_meta_graph(modelbasepath + ".meta")
  saver.restore(sess, modelbasepath)
  ... # restore operations here
  # finally run predictions, error occurs here!
  predictions = sess.run([predictions], feed_dict={handle: training_handle})

我收到此错误:

Number of components does not match: expected 2 types but got 1.
 [[Node: IteratorFromStringHandle_2 = IteratorFromStringHandle[output_shapes=[[?,80,80,3], [?,80,80,?]], output_types=[DT_FLOAT, DT_FLOAT], _device="/job:localhost/replica:0/task:0/device:CPU:0"](_arg_Placeholder_1_0_0)]]

哪个告诉我,我也希望得到一个标签,而我只提供图像来进行预测。

我该如何克服?有没有一种方法可以改变占位符的形状,或者将如何实现,从而可以恢复使用datatset API和可喂养的dict训练的模型?

1 个答案:

答案 0 :(得分:1)

我遇到了同样的问题。但是,我无法提出一个干净的解决方案。我最终为加载图像时返回的标签创建了一个虚拟张量。也许有更好的方法,但是此解决方案现在应允许您运行模型。

dataset = tf.data.Dataset.from_tensor_slices(filenames)
dataset = dataset.map(load_images)

def load_images(x):
    image = tf.image.decode_png(tf.read_file(x), channels = 3))
    image = self.normalize(image)
    image = tf.image.resize_images(image, [IM_WIDTH, IM_HEIGHT])

    # Assuming label is one channel, can slice image to get correct dims
    label = tf.zeros_like(image[:, :, 0:1]) 

    return image, label