@tf.function 正在减慢训练步骤

Question

我正在使用以下 tf.function 修饰训练步骤：

@tf.function
def train_step(inputs, labels):
    with tf.GradientTape(persistent=True) as tape:
        predictions = model([X, F], training=True)
        losses = [l_f(tf.expand_dims(labels[:,i], axis=-1), predictions[i]) for i, l_f in enumerate(loss_functions)]
    gradients = [tape.gradient(l, model.trainable_variables) for l in losses]
    for g in gradients:
        grads = [gg if gg is not None else tf.zeros_like(model.trainable_variables[i], dtype=tf.float32) for i, gg in enumerate(g)]
        optimizer.apply_gradients(zip(grads, model.trainable_variables)
    del tape
    return losses


def weighted_loss(weights):
    @tf.function
    def loss_func(labels, predictions):
        min_class_filter = tfk.backend.greater(labels, 0.5)

        y_min = tf.boolean_mask(labels, min_class_filter)
        y_max = tf.boolean_mask(labels, tf.math.logical_not(min_class_filter))
        y_pred_min = tf.boolean_mask(predictions, min_class_filter)
        y_pred_max = tf.boolean_mask(predictions, tf.math.logical_not(min_class_filter))

        loss_min_class = tfk.backend.mean(tfk.backend.binary_crossentropy(y_min, y_pred_min))
        loss_max_class = tfk.backend.mean(tfk.backend.binary_crossentropy(y_max, y_pred_max))
        loss_all = tfk.backend.mean(tfk.backend.binary_crossentropy(labels, predictions))
        return weights[0]*loss_min_class + weights[1]*loss_max_class + weights[2]*loss_all
    return loss_func

loss_functions = [weighted_loss(w) for w in target_weights]

这有点古怪，但基本上，我的网络有多个输出，这意味着在某些情况下，为某些权重返回 None 的梯度是正确的，所以我将这些梯度替换为零，我正在计算分别在这些输出中的每一个上损失，然后在每一步传播它们中的每一个。

当我按照书面方式运行时，运行单个训练步骤需要很长时间（10 分钟以上），并且我在日志中看到以下消息：

E tensorflow/core/grappler/optimizers/meta_optimizer.cc:502] function_operator failed: Invalid argument: Input 0 of node model/LSTM_forward_0/zeros_like was passed int32 from model/LSTM_forward_0/StatefulPartitioned Call:9 incompatible with expected variant.

当我删除 @tf.function 装饰器时，它在大约 10% 的时间内运行，并且我没有看到此日志警告。这个警告是转移注意力还是合法地指出了通过添加 @tf.function 产生的问题？

其他细节：

• TF 2.0
• GPU 已启用且可用
• CUDA 10.1
• GPU 利用率在这两种情况下都是 0%，但这不是由于数据馈送最大化 CPU 吞吐量引起的，因为当我在训练循环之外生成训练数据时，它与来自 TFRecords 的瞬时值一样好，具有足够的预取和有限的增强
• dtype 的输入、标签、梯度和所有 model.trainable_variables 都是 tf.float32

Answer 1

根据我的阅读，tf.function 不应该包含对图形变量的任何分配，以使其顺利运行。

在训练步骤中，您正在更改模型的权重，因此违反了这一点。

我不确定这是不是这个原因，但是您可以尝试将tf.function 仅留在损失函数中，而不是在训练步骤中。

Answer 2

我已经想出了如何解决它。问题在于覆盖无渐变，而不是持久渐变磁带。

@tf.function
def train_step(inputs, labels):
    with tf.GradientTape(persistent=True) as tape:
        predictions = model([X, F], training=True)
        losses = [l_f(labels, predictions, i) for i, l_f in enumerate(loss_functions)]
    gradients = [tape.gradient(l, model.trainable_variables) for l in losses]
    for g in gradients:
        optimizer.apply_gradients(zip(g, model.trainable_variables)
    del tape
    return losses


def weighted_loss(weights):
    @tf.function
    def loss_func(labs, preds, i):
        labels = tf.expand_dims(labs[:,i], axis=-1)
        predictions = preds[i]
        min_class_filter = tfk.backend.greater(labels, 0.5)

        y_min = tf.boolean_mask(labels, min_class_filter)
        y_max = tf.boolean_mask(labels, tf.math.logical_not(min_class_filter))
        y_pred_min = tf.boolean_mask(predictions, min_class_filter)
        y_pred_max = tf.boolean_mask(predictions, tf.math.logical_not(min_class_filter))

        loss_min_class = tfk.backend.mean(tfk.backend.binary_crossentropy(y_min, y_pred_min))
        loss_max_class = tfk.backend.mean(tfk.backend.binary_crossentropy(y_max, y_pred_max))
        loss_all = tfk.backend.mean(tfk.backend.binary_crossentropy(labels, predictions))
        return weights[0]*loss_min_class + weights[1]*loss_max_class + weights[2]*loss_all
    return loss_func

loss_functions = [weighted_loss(w) for w in target_weights]

通过将所有输出和所有标签传递给损失函数（即使我忽略了其中的一堆），磁带将为所有分支返回适当的梯度 (0)，而不仅仅是针对特定损失的焦点。