在 TF op 中评估 TF 模型会引发错误

Question

我正在使用 TensorFlow 2。我正在尝试优化一个函数，该函数使用经过训练的 tensorflow 模型的损失（毒药）。

@tf.function
def totalloss(x):
    xt = tf.multiply(x, (1.0 - m)) + tf.multiply(m, d)
    label = targetlabel*np.ones(xt.shape[0])
    loss1 = poison.evaluate(xt, label, steps=1)
    loss2 = tf.linalg.norm(m, 1)
    return loss1 + loss2

我无法执行此功能，但是，当我评论 @tf.function 行时，该功能有效！

我需要将此函数用作张量流操作，以优化'm'和'd'。

值错误：未知图表。中止。

这就是我定义模型和变量的方式：

# mask
m = tf.Variable(tf.zeros(shape=(1, 784)), name="m")
d = tf.Variable(tf.zeros(shape=(1, 784)), name="d")
# target
targetlabel = 6
poison = fcn()
poison.load_weights("MNISTP.h5")
adam = tf.keras.optimizers.Adam(lr=.002, decay=1e-6)
poison.compile(optimizer=adam, loss=tf.losses.sparse_categorical_crossentropy)

这就是我稍后调用该函数的方式：（执行此行会导致下面列出的错误。但是，如果我注释掉@tf.function 行，则此命令有效！）

loss = totalloss(ptestdata)

这是整个回溯调用：

ValueError: in converted code:

    <ipython-input-52-4841ad87022f>:5 totalloss  *
        loss1 = poison.evaluate(xt, label, steps=1)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:746 evaluate
        use_multiprocessing=use_multiprocessing)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training_arrays.py:693 evaluate
        callbacks=callbacks)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training_arrays.py:187 model_iteration
        f = _make_execution_function(model, mode)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training_arrays.py:555 _make_execution_function
        return model._make_execution_function(mode)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:2034 _make_execution_function
        self._make_test_function()
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:2010 _make_test_function
        **self._function_kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/backend.py:3544 function
        return EagerExecutionFunction(inputs, outputs, updates=updates, name=name)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/backend.py:3429 __init__
        raise ValueError('Unknown graph. Aborting.')

    ValueError: Unknown graph. Aborting. 

Answer 1

@tf.function装饰器的目的是将用Python编写的Tensorflow操作转换成Tensorflow图，以达到更好的性能。当您尝试使用带有序列化图的预训练模型时，可能会出现错误。因此，装饰器无法进行图到图的转换。

我在这里报告了这个错误：https://github.com/tensorflow/tensorflow/issues/33997

一个（临时）解决方案是您的损失函数应该分成两个小函数。装饰器只能用于不包括预训练模型的函数中。这样，你仍然可以在其他操作中获得更好的性能，但不是使用预训练模型的部分。

例如：

@tf.function
def _other_ops(x):
    xt = tf.multiply(x, (1.0 - m)) + tf.multiply(m, d)
    label = targetlabel * np.ones(xt.shape[0])
    loss2 = tf.linalg.norm(m, 1)

    return xt, label, loss2

def total_loss(x):
    xt, label, loss2 = _other_ops(x)
    loss1 = poison.evaluate(xt, label, steps=1)

    return loss1 + loss2

---

更新：

根据上面TF问题链接中的讨论，一个优雅的解决方案是手动将输入通过模型的每一层。您可以通过调用your_model.layers

获取模型中的层列表

在您的情况下，您可能会根据使用最后一层中的标签对输出的预测来计算损失。因此，我认为您应该跳过最后一层并计算循环外的损失：

@tf.function
def totalloss(x):
    xt = tf.multiply(x, (1.0 - m)) + tf.multiply(m, d)
    label = targetlabel*np.ones(xt.shape[0])

    feat = xt
    # Skip the last layer which calculates loss1
    for i in range(len(poison.layers) - 1):
        layer = poison.layers[i]
        feat = layer(feat)

    # Now, calculate loss by yourself
    loss1 = tf.keras.losses.sparse_categorical_crossentropy(feat, label)
    loss2 = tf.linalg.norm(m, 1)
    return loss1 + loss2

TF 工程师对此问题的解释方式是，模型可能包含由@tf.function 保证的高级处理。因此，不建议将模型放在以@tf.function 装饰的函数中。因此，我们需要将模型分解成更小的部分来绕过它。