自定义损失函数显着减慢了多 GPU 模型的训练

Question

我正在使用 Keras 的 multi_gpu_model 训练蛋白质结构的变分自动编码器。 当从普通 AE 切换到 VAE 时，我的模型每个 epoch 需要超过 3 倍的时间来训练。

我发现问题出在损失函数上，将其改回内置的 mse 结果与之前看到的速度相同。

我正在使用或多或少与许多教程中看到的相同的 vae_loss 实现：

def vae_loss(y_true, y_pred):
    reconstruction_loss = recon_loss(y_true, y_pred)
    kl_loss = beta * K.mean(
             1 + K.flatten(z_log_var) - K.square(K.flatten(z_mean)) - K.exp(K.flatten(z_log_var)), axis=-1)
    kl_loss /= kl_loss_scaling # divide kl_loss by size of output dimension
    total_loss = K.mean(reconstruction_loss + kl_loss)
    return total_loss

在监控 GPU 使用情况时，我意识到它们得到了很好的利用，然后在每个 epoch 后使用率下降到零。 批量大小根据 GPU 数量和完全相同的设置进行调整，但使用 mse 作为损失工作正常。 似乎 GPU 正在等待计算损失，因此有相当长的停机时间。 （对于较小的批次大小，效果更明显，因此增加此参数在某种程度上是一种解决方案，但我认为这远非最佳）。

这是不可避免的，因为这种损失的计算成本更高，还是我可以调整以获得更好的性能？

Answer 1

根本原因可能是 Keras 中自定义损失函数的支持。 如果你在 Keras 中使用预定义的损失，它工作得很好。 您可以尝试的一件事是将损失函数重写为 Lambda 层并将模型更改为多输出，一个是您的原始输出，一个是模型损失。