自定义丢失函数大大降低了多gpu模型的训练速度

Question

我在用Keras的多gpu模型训练蛋白质结构的可变自动编码器。从普通AEs切换到VAE时，我的模型每历元的训练时间要长3倍以上。你知道吗

我将问题确定为损失函数，将其更改回内置的mse将以与之前相同的速度产生结果。你知道吗

我使用的vae\u 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使用情况时，我意识到它们得到了很好的利用，然后在每个历元之后利用率降到零。batchsize根据GPU的数量和完全相同的设置进行调整，但使用mse作为损耗可以正常工作。似乎GPU正在等待计算损失，因此有相当多的停机时间。（对于较小的batchSize，效果更明显，因此增加此参数在某种程度上是一种解决方案，但我认为这远不是最佳的）。你知道吗

这是不可避免的，因为这种损失是更昂贵的计算或有什么我可以调整，以实现更好的性能？你知道吗

Answer 1

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