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

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

Answer 1

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