自定义损失功能会大大减慢多GPU模型的训练

Question

我正在使用Keras的multi_gpu_model训练蛋白质结构上的变型自动编码器。从普通AE切换到VAE时，我的模型每历时训练时间要长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图层，并将模型更改为多个输出，一是您的原始输出，一是模型损失。