Tensorflow 和 Keras：差异输出形状取决于训练或推断答案

【问题标题】：Tensorflow & Keras: Different output shape depends on training or inferringTensorflow 和 Keras：差异输出形状取决于训练或推断
【发布时间】：2020-05-14 04:00:31
【问题描述】：

我正在对tensorflow.keras.Model 进行子类化以实现某个模型。预期行为：

训练（拟合）时间：返回张量列表，包括最终输出和辅助输出；
推断（预测）时间：返回单个输出张量。

代码是：

class SomeModel(tensorflow.keras.Model):
    # ......
    def call(self, x, training=True):
        # ......
        return [aux1, aux2, net] if training else net

这就是我使用它的方式：

model=SomeModel(...)
model.compile(...,
    loss=keras.losses.SparseCategoricalCrossentropy(),
    loss_weights=[0.4, 0.4, 1],...)
# ......
model.fit(data, [labels, labels, labels])

得到：

AssertionError：在转换后的代码中：

ipython-input-33-862e679ab098:140 调用 *
   `return [aux1, aux2, net] if training else net`
...\tensorflow_core\python\autograph\operators\control_flow.py:918 if_stmt

那么问题是if语句被转换为计算图，这当然会导致问题。我发现整个堆栈跟踪很长而且没用，所以这里没有包含。

那么，有没有办法让TensorFlow根据training生成不同的图呢？

【问题讨论】：

标签： python tensorflow keras loss-function

【解决方案1】：

您使用的是哪个 TensorFlow 版本？您可以覆盖 Tensorflow 2.2 中 .fit、.predict 和 .evaluate 方法中的行为，这将为这些方法（我假设）生成不同的图表，并可能适用于您的用例。

早期版本的问题是子类模型是通过跟踪call 方法创建的。这意味着 Python 条件变成了 TensorFlow 条件，并且在图创建和执行过程中面临一些限制。
首先，必须定义两个分支（if-else），并且关于 python 集合（例如列表），分支必须具有相同的结构（例如元素数量）。您可以阅读有关 Autograph here 和 here 的限制和影响。

（此外，如果条件基于 Python 变量而不是张量，则可能不会在每次运行时评估条件。）

【讨论】：

谢谢。最终使自定义模型不扩展 keras.Model 并在调用后返回 keras.Model 函数。也许子类化keras.layers.Layer 更好，我没试过。
我认为子类化 keras.layers.Layer 不会解决问题，因为同样的行为也会发生（call 方法被 Autograph 追踪）。 training 参数应该（我认为）只改变层的行为，而不是结构/类型/等。它的输出。也许我会尝试多个输出头（Functional API），当训练完成后，定义一个只有预测头的新模型并复制权重。
好吧，Layer 有一个额外的 build 方法，所以我认为它可能会有所不同。顺便说一句，如果创建的Layer 的输入形状没有改变，则不需要复制权重。一旦我重新连接了一个经过训练的演示模型并且准确性没有下降，这验证了这一事实。