自定义损失函数大大减慢了多 GPU 模型的训练答案

【问题标题】：Custom loss function slows training of multi gpu model considerably自定义损失函数大大减慢了多 GPU 模型的训练
【发布时间】：2019-12-08 16:50:22
【问题描述】：

我正在使用 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 正在等待计算损失，因此有相当长的停机时间。（对于较小的批量，效果更明显，因此增加此参数在某种程度上是一种解决方案，但我认为这远非最佳）。

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

【问题讨论】：

标签： python tensorflow keras multi-gpu

【解决方案1】：

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

【讨论】：