自定義損失函數顯着減慢了多 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 層並將模型更改為多輸出，一個是您的原始輸出，一個是模型損失。