tensorflow - 将 keras 应用程序模型与 tf.functions 结合使用时出现 "Unknown graph"错误

Question

这是我的代码:

import tensorflow as tf
import tensorflow_datasets as tfds
import tensorflow.keras.applications.vgg16 as vgg16

tf.enable_eager_execution()

def resize_image(image, shape = (224,224)):
  target_width = shape[0]
  target_height = shape[1]
  initial_width = tf.shape(image)[0]
  initial_height = tf.shape(image)[1]
  im = image
  ratio = 0
  if(initial_width < initial_height):
    ratio = tf.cast(256 / initial_width, tf.float32)
    h = tf.cast(initial_height, tf.float32) * ratio
    im = tf.image.resize(im, (256, h), method="bicubic")
  else:
    ratio = tf.cast(256 / initial_height, tf.float32)
    w = tf.cast(initial_width, tf.float32) * ratio
    im = tf.image.resize(im, (w, 256), method="bicubic")
  width = tf.shape(im)[0]
  height = tf.shape(im)[1]
  startx = width//2 - (target_width//2)
  starty = height//2 - (target_height//2)
  im = tf.image.crop_to_bounding_box(im, startx, starty, target_width, target_height)
  return im

def scale16(image, label):
  im = resize_image(image)
  im = vgg16.preprocess_input(im)
  return (im, label)

data_dir = "/content/"
model = vgg16.VGG16(weights='imagenet', include_top=True)

datasets, info = tfds.load(name="imagenet2012", 
                            with_info=True, 
                            as_supervised=True, 
                            download=False, 
                            data_dir=data_dir, shuffle_files=False
                            )
train = datasets['train'].map(scale16).batch(2).take(1)

@tf.function
def predict_batch(b, m):
  return m.predict_on_batch(b)

sample_imgs = train
for x in sample_imgs:
  print("batch prediction", predict_batch(x[0], model))

当我在没有@tf.function的情况下运行代码时，我得到了预期的结果(批处理的预测)。当我使用 @tf.function 运行它时，出现此错误:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-1-ba7af933ed07> in <module>()
     49 sample_imgs = train
     50 for x in sample_imgs:
---> 51   print("batch prediction", predict_batch(x[0], model))

7 frames
/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/framework/func_graph.py in wrapper(*args, **kwargs)
    903           except Exception as e:  # pylint:disable=broad-except
    904             if hasattr(e, "ag_error_metadata"):
--> 905               raise e.ag_error_metadata.to_exception(e)
    906             else:
    907               raise

ValueError: in converted code:

    <ipython-input-1-ba7af933ed07>:47 predict_batch  *
        return m.predict_on_batch(b)
    /usr/local/lib/python3.6/dist-packages/tensorflow_core/python/keras/engine/training.py:1155 predict_on_batch
        self._make_predict_function()
    /usr/local/lib/python3.6/dist-packages/tensorflow_core/python/keras/engine/training.py:2179 _make_predict_function
        **kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow_core/python/keras/backend.py:3669 function
        return EagerExecutionFunction(inputs, outputs, updates=updates, name=name)
    /usr/local/lib/python3.6/dist-packages/tensorflow_core/python/keras/backend.py:3553 __init__
        raise ValueError('Unknown graph. Aborting.')

    ValueError: Unknown graph. Aborting.

我正在 Google Colaboratory 上运行代码。为什么我会收到此错误？训练步骤通常不应该在 @tf.function 中吗？为什么我不能在这个方法中使用模型？

Answer 1

面临同样的问题。我已经在这里报告了:https://github.com/tensorflow/tensorflow/issues/33997